알고리즘 :: 이것이 코딩 테스트다 :: 이진탐색 :: Q27 :: 정렬된 배열에서 특정 수의 개수 구하기

문제

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있을 때 수열에서 x가 등장하는 횟수를 계산하시오. 단, O(logN) 알고리즘으로 풀어야 한다.

문제접근

이번 챕터에서는 이진 탐색의 새로운 기능을 배우게 된다.
이 문제의 제목이기도 한데, 이진 탐색을 이용하면 정렬된 배열에서 특정 원소의 개수를 구할 수 있다.

• 정렬된 배열과 O(logN)이라는 힌트를 기반으로 이 문제는 이진 탐색 문제임을 파악한다.
• C++에서는 특정 원소의 개수를 구할 때 아주 편리한 함수가 있다.
  • algorithm 헤더파일에 있는 lower_bound() 와 upper_bound() 함수다.
• 따라서 이 문제는 두 가지 방법으로 해결할 수 있다.
  1. 이진 탐색을 구현해서 해결하는 방법.
  2. STL을 사용해서 해결하는 방법,

코드

1. 이진 탐색을 구현해서 해결하는 방법

#include <iostream>
using namespace std;

static int N, x, arr[1000000];

int getFirstPos(int left, int right) {
    if (left > right) return 0;
    int mid = (left + right) / 2;

    // `mid`가 시작 원소거나, mid는 x인데 mid - 1이 x보다 작다면 x의 가장 좌측 원소이다.
    if ((mid == 0 || x > arr[mid - 1]) && x == arr[mid]) return mid;
    else if (arr[mid] < x) getFirstPos(mid + 1, right);
    else getFirstPos(left, mid - 1);
}

int getLastPos(int left, int right) {
    if (left > right) return 0;
    int mid = (left + right) / 2;

    // `mid`가 끝 원소거나, mid는 x인데 mid + 1이 x보다 크다면 x의 가장 우측 원소이다.
    if ((mid == N - 1 || x < arr[mid + 1]) && x == arr[mid]) return mid;
    else if (arr[mid] <= x) getLastPos(mid + 1, right);
    else getLastPos(left, mid - 1);
}

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    cin >> N >> x;   // N개의 원소와 찾으려는 수 x
    for (int i = 0; i < N; ++i) cin >> arr[i];
	
    int answer = getFirstPos(0, N - 1) - getLastPos(0, N - 1);
    if (answer != 0) cout << answer + 1 << '\n';
    else cout << -1 << '\n';
}

• getFirstPos() 함수는 원소 x가 나타나는 시작 지점(인덱스)을 반환하는 함수다.
  • 탐색에 성공하는 경우
    1. mid를 변경하다가 0번 인덱스에 도착하는 경우 (범위 초과 방지)
    2. mid-1번 인덱스에 들어있는 요소는 x보다 작고, mid번 인덱스는 x인 경우
  • mid를 이동하는 경우
    1. mid번 인덱스의 값이 x보다 작다면 오른쪽 부분을 탐색해야 한다.
    2. mid번 인덱스의 값이 x보다 크다면 왼쪽 부분을 탐색해야 한다.
• getLastPos() 함수는 원소 x가 나타나는 종료 지점(인덱스)을 반환하는 함수다.
  • 탐색에 성공하는 경우
    1. mid를 변경하다가 N-1번 인덱스에 도착하는 경우 (범위 초과 방지)
    2. mid+1번 인덱스에 들어있는 요소는 x보다 크고, mid번 인덱스는 x인 경우
  • mid를 이동하는 경우
    1. mid번 인덱스의 값이 x보다 작거나 같다면 오른쪽 부분을 탐색해야 한다.
    2. mid번 인덱스의 값이 x보다 크다면 왼쪽 부분을 탐색해야 한다.

2. STL을 사용해서 해결하는 방법

#include <iostream>
#include <algorithm>
using namespace std;

static int N, x, arr[1000000];

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    cin >> N >> x;
    for (int i = 0; i < N; ++i) cin >> arr[i];
	
    int answer = upper_bound(arr, arr + N, x) - lower_bound(arr, arr + N, x);
    if (answer != 0) cout << answer << '\n';
    else cout << -1 << '\n';
}