Binary Search - 정렬된 배열에서 특정 수의 개수 구하기

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

입력 예시
7 2
1 1 2 2 2 2 3

출력 예시
4
"""
import bisect

"""
해당 문제를 푸는데 기억해야할 것은 '정렬'된 배열이라는 것이다.
즉, target의 숫자의 가장 왼쪽과 오른쪽 idx를 구한 후 해당 idx의 차이를 반환하면 된다.
"""
def solution(lst, target):
    left_idx = bisect.bisect_left(lst, target)
    if not (0 <= left_idx < len(lst) and lst[left_idx] == target):
        return -1

    right_idx = bisect.bisect_right(lst, target)
    return right_idx - left_idx


assert solution(lst=[1,1,2,2,2,2,3], target=2) == 4

특정 값의 인덱스를 bisect.bisect_left 모듈 함수를 사용하여 알아낸 후 bisect.bisect_right 모듈 함수를 이용하여 특정 값의 끝 이후 인덱스를 알아내어 두 값 사이의 차이를 도출하여 갯수를 알아낸다.

 

단, 처음 left_idx를 구할 때 해당 인덱스는 주어진 배열의 길이에 속해야 하며, 도출한 idx가 lst[idx]와 target 이 같은지 조건을 건다.