[Zoho인터뷰] 정렬된 배열에서 특정 수의 개수 구하기 (python)

문제

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

 

 

풀이

이 문제는 입력 조건에서 N이 최대 100만이 될 수 있기 때문에 시간복잡도 O(logN)을 가진 알고리즘을 이용해야 한다. 이렇게 입력되는 데이터의 값이 클 때에는 이진탐색을 고려해볼 수 있다. 

입력되는 데이터가 정렬되어 있다고 했으므로 같은 수가 있을 때에는 연속되어 나열되어 있을 것이다. 

처음에 나는 이진탐색을 통해 특정 수를 찾았으면, 양쪽으로 이동하면서 다른 수가 나올 때까지 반복문을 돌렸다. 하지만 이럴 경우 100만개의 데이터가 주어지고 그 중 80만 개가 찾고 있는 타겟이 된다고 하면 80만 개의 데이터에 대해서 반복문으로 하나씩 방문해야 하는 알고리즘이므로 매우 비효율적이다.

 

그래서 이진탐색으로 특정 수의 개수를 구하는 방법은 바로

특정 수가 시작하는 인덱스와, 마지막으로 나오는 인덱스를 구해 그 차이를 구하는 것이다. 이진탐색을 2번 거쳐 이런 방법으로 특정 수의 개수를 구할 수 있다. 

 

import sys
input = sys.stdin.readline
n, x = map(int, input().rstrip().split())
arr = list(map(int, input().split()))

arr.sort()

count = 0
def first(arr, start, end, target):
    if start > end:
        return -1

    mid = (start+end)//2
    if (mid == 0 or arr[mid-1] < target) and (arr[mid] == target):
        return mid
    elif arr[mid] >= target:
        return first(arr, start, mid-1, target)
    else:
        return first(arr, mid+1, end, target)

def last(arr, start, end, target):
    if start > end:
        return None
    mid = (start + end)//2
    if (mid == n-1 or arr[mid+1] > target) and (arr[mid] == target):
        return mid
    elif arr[mid] > target:
        return last(arr, start, mid - 1, target)
    else:
        return last(arr, mid+1, end, target)

def count_target(arr, x):
    a = first(arr, 0, n-1, x)
    # 수열에 x가 존재하지 않는 경우
    if a == -1:
        return 0
    b = last(arr, 0,n-1,x)
    cnt = (b-a) + 1
    return cnt

result = count_target(arr, x)
print(result)

 

기본적인 이진탐색 알고리즘에서는 타겟 넘버 하나만 찾을 수 있는데,

만약 수열에 같은 수가 여러번 나올 경우, 이진탐색 2개를 이용해 여러번 나오는 수의 첫번째 인덱스와 마지막 인덱스를 찾을 수 있다. 이 방식은 어려울 이진탐색 문제에서 또 나올 수 있을 것 같아서 기억해둬야 겠다.

 

알고리즘 분류

• 이진탐색