[이코테] 이진탐색 - 정렬된 배열에서 특정 수의 개수 구하기

해당 포스팅의 문제의 출처는 나동빈님의 이것이 취업을 위한 코딩 테스트 교재를 공부하면서 풀이할 때 본인의 사고 과정 흐름과 문제 풀이를 기록하기 위함 입니다. 

 

---

문제설명

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이 때 이 수열에서 x가 등장하는 횟수를 계산하여라. 예를 들어 수열 [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을 출력한다.

사고과정

• 이미 정렬되었다는 배열을 확인한 후 이진탐색을 재귀함수를 활용하여 구현하려 했지만 원소를 탐색까지는 하되 그 원소 개수를 세는 아이디어를 시간 내에 떠올리지 못했다.
• 풀이를 보았고, 핵심적인 포인트는 원소는 이미 모두 정렬되어 있기 때문에 수열 내에 x가 존재한다면 무조건 연속적으로 나열되어 있을 것이라는 점. 그렇기 때문에 연속적으로 나열된 x의 가장 왼쪽 인덱스 번호와 가장 오른쪽 인덱스 번호 차이값 + 1을 하게 되면 x 원소의 개수를 출력할 수 있다!
• 이진 탐색을 손쉽게 구현할 수 있는 bisect 라이브러리 관련 함수들에 대한 존재를 몰랐었다. 이를 알았다면 매우 손쉽게 해결할 수 있다..

풀이

1. Native하게 Python으로 구현하는 풀이

import sys

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

# 1.포인트: 이진탐색을 활용 
# 2.포인트: 정렬되어있으므로 원소 개수 세기 위해서 탐색하려는 값이 가장 왼쪽에있는 인덱스 & 가장 오른쪽 인덱스 차이 + 1 하면 원소의 개수임


# 가장 왼쪽에 있는 타겟값 인덱스 찾는 함수
def first(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 가장 왼쪽에 있는 값을 발견했을 경우
    if (mid == 0 or array[mid-1] < target) and array[mid] == target:
        return mid
    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 array[mid+1] > target) 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)


def count_x_value(array, target):
    N = len(array)
    start = 0
    end = N-1

    first_idx = first(array, target, start, end)
    if first_idx == None:
        return 0
    last_idx = last(array, target, start, end)

    return last_idx - first_idx + 1


res = count_x_value(array, x)
if res == 0:
    print(-1)
else:
    print(res)

2. bisect 라이브러리 활용

from bisect import bisect_left, bisect_right

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


def count_x_range(array, left_val, right_val):
    left_idx = bisect_left(array, left_val)
    right_idx = bisect_right(array, right_val)
    print(left_idx, right_idx)
    return right_idx - left_idx

res = count_x_range(array, x, x)
# 똑같은 원소값들을 찾는 것이기 때문에 해당 원소가 없으면 무조건 left_idx, right_idx 동일한 값으로 나오게 됨!
if res == 0:
    print(-1)
else:
    print(res)