Binary Search - 정렬된 배열에서 수의 개수 구하기 (이분탐색)

정렬된 숫자의 배열에서 특정 수의 개수를 구한다.
시간복잡도가 O(logN) 이하여야 시간초과 판정을 받지 않는다.
수가 배열에 존재하지 않는 경우 -1을 출력한다.

입력
n: 배열의 길이(1~1000000), x: 찾을 수

Solution

import bisect

n,x=map(int,input().split())
num = sorted(list(map(int, input().split())))
count=bisect.bisect_right(num,x)-bisect.bisect_left(num,x)
if count<0: print(-1)
else: print(count)

특정 수의 개수는 count() 메소드로 찾을 수 있지만, 정렬된 배열의 조건, 배열의 길이를 봐서 일반적인 loop로 완전탐색으로 돌다가는 당연히 시간초과가 날 것을 알고 이분 탐색이라는 포인트를 잡을 수 있어야 한다.

bisect 라이브러리에서 bisect_right() 메소드와 bisect_left() 메소드를 사용하면 첫번째 인자의 배열에 두번째 요소를 집어넣을 때 각각 오른쪽에서 넣을 경우, 왼쪽에서 넣을 경우에서 몇 번째 인덱스에 넣어야 하는지를 반환한다.

예를 들어, 배열이 1 1 2 2 2 2 3으로 구성되었다고 가정해보자.
2의 개수를 찾는 경우 2를 오른쪽부터 넣을 때의 인덱스는 6이고, 왼쪽부터 넣을 때의 인덱스는 2이다. 이 두 값을 빼면 바로 개수가 된다. 2가 없을 경우 두 값이 모두 같을 것이고, 자연스레 count는 0이 된다.

이분탐색 방식으로 코드를 직접 구현해도 되지만 외부 라이브러리 제한조건이 없는 경우 번거로운 수고를 덜수도 있다. 하지만 라이브러리는 기초를 알고 쓰자.