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

문제

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요. 단, 이 문제는 시간 복잡도 O (log N)으로 알고리즘을 설계하지 않으면 시간 초과 판정을 받습니다.

입력 조건 : 첫째 줄에 N과 x가 정수형태로 공백으로 구분되어 있고, 둘째 줄에 N개의 원소가 정수 형태로 구분되어 입력된다.
출력 조건 : 수열의 원소 중에서 값이 x인 원소의 개수를 출력하고, x인 원소가 하나도 없다면 -1을 출력한다.

https://www.youtube.com/watch?v=94RC-DsGMLo&list=PLRx0vPvlEmdAghTr5mXQxGpHjWqSz0dgC&index=5

22:35에 나오는 문제입니다.

---

만약, 이 문제의 제한조건에 시간복잡도가 O(log N)이라는 조건이 없었다면, 단순이 리스트를 선형탐색해서 x면 count += 1을 했을 테지만, 이는 O(N)이기 때문에 안됩니다!

하지만 강의 앞부분에서  파이썬 이진탐색 라이브러리인 bisect_left와 bisect_right를 통해 인덱스를 찾거나 개수를 세는 걸 알려주어서 쉽게 코드를 구현할 수 있었습니다.

bisect_left(array, x) 함수는 이미 정렬되어있는 배열 array에서 가장 왼쪽부터(0번 인덱스부터) x를 찾아나가다 가장 먼저 발견한 x의 인덱스를 반환합니다.
bisect_right(array, x) 함수는 이미 정렬되어있는 배열 array에서 가장 오른쪽부터 x를 찾아나가다 가장 먼저 발견한 x의 다음 인덱스를 반환합니다. 설명을 쉽게 하기 위해 x의 인덱스를 반환한다고 했지만, 사실 x가 들어갈 수 있는 인덱스를 반환하는 것입니다.
예를 들어, [1, 2, 2, 3] 과 같은 정렬되어 있는 배열이 있을 때, bisect_left(array, 2)를 하면 왼쪽부터 탐색해 나가며 2가 들어갈 수 있는 가장 왼쪽 인덱스를 반환합니다. 그럼 [1, 여기 , 2, 2, 3] 에 들어갈 수 있는 것이기 때문에 1을 반환합니다.
반대로 bisect_right(array, 2)를 하게 되면, 오른쪽부터 탐색해 나가며 2가 들어갈 수 있는 가장 오른쪽 인덱스를 반환합니다. 그럼 [1, 2, 2, 여기, 3]에 들어갈 수 있는 것이기 때문에 3을 반환합니다.

그럼 [1, 2, 3, 3, 3, 4] 와 같이 정렬된 배열 a가 있을 때 bisect_left(a, 3)을 하면 2가 출력되고, bisect_right(a, 3)을 하면 5가 출력되게 되겠죠? bisect_right(a, 3)를 했을 때 4가 아닌 5 임에 주의하셔야 합니다!

이를 활용하여 배열에 있는 x의 개수를 빠르게 알아낼 수도 있습니다.
bisect_right(a, 3)  - bisect_left(a, 3)와 같이 같은 수 x를 오른쪽, 왼쪽에서 찾아서 빼면 개수가 나옵니다!

이를 활용하여 다음과 같이 코드를 작성했습니다.

from bisect import bisect_left, bisect_right

def count_by_range(a, left_value, right_value):
  right_index = bisect_right(a, right_value)
  left_index = bisect_left(a, left_value)
  return right_index-left_index

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

count = count_by_range(num_list,x,x)

print(-1 if count==0 else count_by_range(num_list,x,x))

여기서 print문 안에 조건문을 써서 더 간결히 적어봤습니다.

print문 안에 조건문을 사용하는 방식은 다음과 같습니다.

print(참 if 조건 else 거짓)