[Python] 정렬된 배열에서 특정 수의 개수 구하기

문제

• N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다.
• 이때 이 수열에서 x가 등장하는 횟수를 계산하세요.
• 예를 들어 수열 {1, 1, 2 ,2 ,2 ,2, 3}이 있을 때 x=2라면, 현재 수열에서 값이 2인 원소가 4개이므로 4를 출력합니다.
• 단, 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 ‘시간 초과’ 판정을 받습니다.

입력

• 첫째 줄에 N과 x가 정수 형태로 공백으로 구분되어 입력됩니다.
  • 1 <= N <= 1,000,000
  • 1e9 <= x <= 1e9

출력

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

입력 예시

7 2
1 1 2 2 2 2 3

출력 예시

4

문제 풀이

• 타겟 값이 나오는 첫 번째 인덱스와 마지막 인덱스의 차를 구하면 타겟 값이 몇 번 등장하는 지 알 수 있다
• Bisect 라이브러리를 사용하면 각 값의 제일 왼쪽 인덱스 및 오른쪽 인덱스 쉽게 찾을 수 있다

코드

• 이진탐색 구현 코드

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

# 정렬된 수열에서 값이 x인 원소의 개수를 세는 함수
def count_by_value(array, x):
    # 데이터의 개수
    n = len(array)
    first = find_first(array, x, 0, n - 1)
    if first == None:
        return 0
    last = find_last(array, x, 0, n - 1)
    return last - first + 1

def find_first(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 가장 왼족에 타겟 값이 있을 때 인덱스 반환
    if array[mid] == target and (mid == 0 or array[mid - 1] < target):
        return mid
    # 타겟 값과 같으나 제일 왼쪽이 아니거나, 타겟 값 보다 현재 값이 더 클 때 탐색 범위 줄이기
    elif array[mid] >= target:
        return find_first(array, target, start, mid - 1)
    # 타겟 값보다 작을 때
    else:
        return find_first(array, target, mid + 1, end)

def find_last(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 가장 오른쪽에 타겟 값이 있을 때 인덱스 반환
    if array[mid] == target and (mid == n - 1 or array[mid + 1] > target):
        return mid
    # 타겟 값 보다 현재 값이 더 클 때 탐색 범위 줄이기
    elif array[mid] > target:
        return find_last(array, target, start, mid - 1)
    # 타겟 값보다 작을 때
    else:
        return find_last(array, target, mid + 1, end)

count = count_by_value(array, x)
if count == 0:
    print(-1)
else:
    print(count)

from bisect import bisect_left, bisect_right
n, x = map(int,input().split())
array = list(map(int,input().split()))

def count_by_range(array, left_value, right_value):
    left_index = bisect_left(array, left_value)
    right_index = bisect_right(array, right_value)
    return right_index - left_index
count = count_by_range(array, x, x)

if count == 0:
    print(-1)
else:
    print(count)