[Binary Search] 이코테 “정렬된 배열에서 특정 수의 개수 구하기” Python 풀이

문제

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

입력 조건

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

출력 조건

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

입력 예시

7 2
1 1 2 2 2 2 3

7 4
1 1 2 2 2 2 3

출력 예시

4

-1

풀이

• 이 문제는 시간 복잡도 logN으로 동작하는 알고리즘을 요구하고 있기 때문에, 일반적인 선형 탐색으로는 문제를 해결할 수 없다.
• 원소들은 모두 정렬되어 있기 때문에, 수열 내에 x가 존재한다면 연속적으로 나열되어 있을 것으로 예상할 수 있다.
• 따라서 x가 처음 등장하는 인덱스와 x가 마지막으로 등장하는 인덱스를 각각 계산한 뒤에, 그 인덱스의 차이를 계산하여 문제를 해결할 수 있다.
  • 핵심 알고리즘이다.
• 해당 알고리즘은 이진 탐색 함수를 2개 작성하여 해결한다.
  • 데이터가 존재한다면 가장 첫 번째 위치를 찾는 이진 탐색 함수
  • 데이터가 존재한다면 가장 마지막 위치를 찾는 이진 탐색 함수
• 단순히 정렬된 배열에서 특정한 데이터를 찾도록 요구하는 문제에서는, 이진 탐색을 직접 구현할 필요 없이, 파이썬의 표준 라이브러리인 bisect 모듈을 사용할 수도 있다.

예제 코드

• 이진 탐색을 직접 구현한 예제 코드

    import sys
  
    # 정렬된 수열에서 값이 x인 원소의 개수를 세는 메서드
    def count_by_value(array, x):
        # 데이터의 개수
        n = len(array)
  
        # x가 처음 등장한 인덱스 계산
        a = first(array, x, 0, n-1)
  
        # 수열에 x가 존재하지 않는 경우
        if a == None:
            return 0 # 값이 X인 원소의 개수는 0개이므로 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:
            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 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 입력 받기
    n, x = map(int, sys.stdin.readline().split())
  
    # 전체 데이터 입력 받기
    array = list(map(int, sys.stdin.readline().split()))
  
    # 값이 X인 데이터의 개수 계산
    count = count_by_value(array, x)
  
    # 값이 x인 원소가 존재하지 않는다면
    if count == 0:
        print(-1)
    # 값이 x인 원소가 존재한다면
    else:
        print(count)

• bisect 모듈 기반의 count_by_range() 함수를 사용한 예제 코드

    import sys
    from bisect import bisect_left, bisect_right
  
    # 값이 [left_value, right_value]인 데이터의 개수를 반환하는 함수
    def count_by_range(array, left_value, right_value):
        right_index = bisect_right(array, right_value)
        left_index = bisect_left(array, left_value)
        return right_index - left_index
  
    # 데이터의 개수 N, 찾고자 하는 값 x 입력 받기
    n, x = map(int, sys.stdin.readline().split())
  
    # 전체 데이터 입력 받기
    array = list(map(int, sys.stdin.readline().split()))
  
    # 값이 [x, x] 범위에 있는 데이터의 개수 계산
    count = count_by_range(array, x, x)
  
    # 값이 x인 원소가 존재하지 않는다면
    if count == 0:
        print(-1)
    # 값이 x인 원소가 존재한다면
    else:
        print(count)

참고

• 이것이 취업을 위한 코딩 테스트다. with python