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

문제

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

입력 조건

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

출력 조건

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

입력 예시 1

7 2
1 1 2 2 2 2 3

출력 예시 1

4

입력 예시 2

7 4
1 1 2 2 2 2 3

출력 예시 2

-1

 

Solution

x가 위치한 첫번째 인덱스와 마지막 인덱스의 차를 이용해서 구한다. 이를 위해 각각의 함수를 만든다.

 

def binary_search(array, target, start, end):
    while start <= end:
        mid = (start + end) // 2
        if array[mid] == target:
            return mid
        elif array[mid] > target:
            end = mid - 1
        else:
            start = mid + 1
    return None

def binary_search_end(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 해당 값을 가지는 원소 중에서 가장 오른쪽에 있는 경우에만 인덱스 반환
    if (mid == n-1 or target < array[mid+1]) and array[mid] == target:
        return mid
    # 중간점의 값 보다 찾고자 하는 값이 작은 경우 왼쪽 확인
    elif array[mid] > target:
        return binary_search_end(array, target, start, mid-1)
    # 중간점의 값 보다 찾고자 하는 값이 크거나 같은 경우 오른쪽 확인
    else:
        return binary_search_end(array, target, mid+1, end)

n, x = map(int, input().split())
data = list(map(int, input().split()))
cnt = 0


if binary_search(data, x, 0, len(data)-1) == None:
    print(-1)
else:
    print(binary_search_end(data, x, 0, len(data)-1) - binary_search(data, x, 0, len(data)-1) + 1)