N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요. 예를 들어 수열 {1,1,2,2,2,2,3}이 있을 때 x = 2라면, 현재 수열에서 값이 2인 원소가 4개이므로 4를 출력합니다.
단, 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 '시간 초과' 판정을 받습니다.
7 2
1 1 2 2 2 2 3
=> 4
7 4
1 1 2 2 2 2 3
=> -1
def count_by_value(array, x):
# 데이터의 개수
n = len(array)
# x가 처음 등장한 인덱스 계산
a = first(array, x, 0, n-1) # array = 졍렬된 배열, x = 찾고자 하는 값, start = 처음 인덱스, end = 종료 인덱스
# 수열에 x가 존재하지 않는 경우
if a == None:
return 0
# x가 마지막으로 등장한 인덱스 계산
b = last(array, x, 0, n-1)
return b - a + 1
# 처음 위치를 찾는 이진 탐색 함수
def first(array, target, start, end):
if start > end:
return None
mid = (start + end) // 2
# 해당 값을 각지는 원소 중에서 가장 왼쪽에 있는 경우에만 인덱스 반환
if (mid == 0 or target > array[mid-1]) and array[mid] == target:
# mid가 배열의 마지막 원소거나, 오른쪽에 target이 없을 때만 반환
# array[mid] == target이면서, 가장 오른쪽 값인지 확인
return mid
# 중간점의 값 보다 찾고자 하는 값이 작거나 같은 경우 왼쪽 확인
# array[mid]이 target보다 크면 → 왼쪽 탐색
elif array[mid] >= target:
return first(array, target, start, mid - 1)
# 중간점의 값 보다 찾고자 하는 값이 큰 경우 오른쪽 확인
else:
return first(array, target, mid+1, end)
# 마지막 위치를 찾는 이진 탐색 함수
def last(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 last(array, target, start, mid-1)
# 중간점의 값 보다 찾고자 하는 값이 크거나 같은 경우 오른쪽 확인
else:
return last(array, target, mid+1, end)
n, x = map(int, input().split())
array = list(map(int, input().split()))
count = count_by_value(array, x)
if count == 0:
print(-1)
else:
print(count)