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

문제

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

입력조건

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

출력조건

• 수열의 원소 중에서 값이 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

풀이 방식

해당 문제는 조건에서부터 O(logN) 알고리즘을 설계해서 풀어야 '시간초과' 판정이 나지 않는다 쓰여있습니다. O(N)의 시간복잡도를 갖는 순차탐색을 사용하지 않고 왜 O(logN)의 시간복잡도를 갖는 이분탐색을 써야할까 생각을 했다. 약 1억번의 연산을 하는 것이 1초로 알고있는데, 해당 문제는 시간 제한이 1초이고, N이 1,000,000이기 때문에 괜찮지 않을까 했다. 그러나 문제에서 O(logN)의 시간복잡도로 풀이를 원하니 이분탐색을 진행하였다. 이 문제는 우선 2번의 이분탐색을 진행해야했다.

x의 값을 갖는 원소에서 가장 작은 index를 찾는 이분탐색과 가장 큰 index를 찾는 이분탐색을 해야했다. 자세한 것은 아래 알고리즘에서 더 설명하겠다.

알고리즘

start와 end를 통해 구하는 mid값은 아래와 같다.

mid = (start + end) // 2

위를 토대로 2번의 이분탐색을 진행하면 된다.

• 가장 작은 index를 찾는 이분탐색
  - array[mid]값이 x보다 크다면 end 값을 더 낮게 변경
  - array[mid]값과 x가 같다면 end 값을 더 낮게 변경해서 같은 원소 있는지 체크
  - array[mid]값이 x보다 작다면 start 값을 더 크게 변경
• 가장 큰 index를 찾는 이분탐색
  - array[mid]값이 x보다 크다면 end 값을 더 낮게 변경
  - array[mid]값과 x가 같다면 start 값을 더 크게 변경해서 같은 원소 있는지 체크
  - array[mid]값이 x보다 작다면 start 값을 더 크게 변경
  

코드

import sys

# 가장 첫번째 x 위치 찾는 함수
def leftBinarySearch(start, end):
  global minIdx

  # 탈출 조건
  if start > end:
    return

  mid = (start + end) // 2  # mid 값 설정

  if array[mid] > x:  # array[mid]값이 x보다 크다면 end 값을 더 낮게 변경
    leftBinarySearch(start, mid - 1)
  elif array[mid] == x:  # array[mid]값과 x가 같다면 end 값을 더 낮게 변경해서 같은 원소 있는지 체크
    minIdx = mid  # minIdx 최신화
    leftBinarySearch(start, mid - 1)
  else:  # array[mid]값이 x보다 작다면 start 값을 더 크게 변경
    leftBinarySearch(mid + 1, end)

# 가장 마지막 x 위치 찾는 함수
def rightBinarySearch(start, end):
  global maxIdx

  # 탈출 조건
  if start > end:
    return

  mid = (start + end) // 2  # mid 값 설정

  if array[mid] > x:  # array[mid]값이 x보다 크다면 end 값을 더 낮게 변경
    rightBinarySearch(start, mid - 1)
  elif array[mid] == x:  # array[mid]값과 x가 같다면 start 값을 더 크게 변경해서 같은 원소 있는지 체크
    maxIdx = mid
    rightBinarySearch(mid + 1, end)
  else:  # array[mid]값이 x보다 작다면 start 값을 더 크게 변경
    rightBinarySearch(mid + 1, end)

N, x = map(int, sys.stdin.readline().split())
array = list(map(int, sys.stdin.readline().split()))
minIdx = N  # 가장 첫번째 x 위치
maxIdx = -1  # 가장 마지막 x 위치

leftBinarySearch(0, N - 1)  # 가장 첫번째 x 위치 찾는 함수
rightBinarySearch(0, N - 1)  # 가장 마지막 x 위치 찾는 함수

if maxIdx == -1 and minIdx == N:  # x인 원소가 존재하지 않았다면 -1 출력
  print(-1)
else:
  print(maxIdx - minIdx + 1)