이진 탐색 - 정렬된 배열에서 특정 수의 개수 구하기

< 정렬된 배열에서 특정 수의 개수 구하기 > 

[문제]

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요. 예를 들어 수열 {1,1,2,2,2,2,3}이 있을 때 x = 2라면, 현재 수열에서 값이 2인 원소가 4개이므로 4를 출력합니다.

단, 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 '시간 초과' 판정을 받습니다.

[입력]

• 첫째 줄에 N과 x가 정수 형태로 공백으로 구분되어 입력됩니다. (1<=N<=1,000,000), (-10e9<=x<10e9)
• 둘째 줄에 N개의 원소가 정수 형태로 공백으로 구분되어 입력됩니다. (-10e9<=각 원소의 값<=10e9)

[출력]

• 수열의 원소 중에서 값이 x인 원소의 개수를 출력합니다. 단, 값이 x인 원소가 하나도 없다면 -1을 출력합니다.

예제 입력	예제 출력
7 2 1 1 2 2 2 2 3	4

---

아이디어

이진 탐색을 통해 배열에 있는 타겟의 맨 앞 인덱스와 맨 뒤 인덱스를 각각 구한다. 그래서 그 길이를 구해주면 된다.

풀이 및 코드

from bisect import bisect_left, bisect_right

n, x = map(int, input().split())
sortedArray = list(map(int, input().split()))

def myBisectLeft(array, target):
  start = 0
  end = n - 1
  idx = -1
  while (start <= end):
    mid = (start + end) // 2
    if sortedArray[mid] == x:
      # 인덱스 저장 후 앞의 배열 다시 탐색
      idx = mid
      end = mid - 1
    elif sortedArray[mid] > x:
      end = mid - 1
    elif sortedArray[mid] < x:
      start = mid + 1
  return idx

def myBisectRight(array, target):
  start = 0
  end = n - 1
  idx = -2
  while (start <= end):
    mid = (start + end) // 2
    if sortedArray[mid] == x:
      # 인덱스 저장 후 뒤의 배열 다시 탐색
      idx = mid
      start = mid + 1
    elif sortedArray[mid] > x:
      end = mid - 1
    elif sortedArray[mid] < x:
      start = mid + 1
  return idx + 1

first = bisect_left(sortedArray, x)
last = bisect_right(sortedArray, x) - 1
print(last - first + 1)

first = myBisectLeft(sortedArray, x)
last = myBisectRight(sortedArray, x) - 1
print(last - first + 1)

bisect 라이브러리에 존재하는 bisect_left, bisect_right함수를 이용해 맨 앞, 맨 뒤 인덱스를 구해준 다음에 그 차이를 출력했다. 하지만 문제의 의도는 이 함수를 이용하라는 게 아니고, 직접 구현해보라는 것 같았다. 그래서 bisect_left, right함수를 직접 구현해보았다.

같은 결과를 출력하였다.

예시를 input으로 주었더니 4를 출력하였다.

후기

이진탐색 굉장히 빠르고 좋은 알고리즘이지만 익숙해지려면 조금 더 공부해야할 것 같다.