[python] 이것이코딩테스트다 : 정렬된 배열에서 특정 수의 개수 구하기 (bisect)

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요.
단, 이 문제의 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 '시간 초과' 판정을 받습니다.

입력 예시

7 2
1 1 2 2 2 2 3

출력 예시

4

 

 

1. 정답 코드  

import sys

def countValue(arr, target):
    first = 0
    last = len(arr)-1

    start = findFirstIndex(arr, target, first, last)
    if start == None:
        return 0
    end = findLastIndex(arr, target, first, last)
    return end - start + 1


# 가장 맨 왼쪽에 있는 원소 인덱스 구하기
def findFirstIndex(arr, target, first, last):
    if first > last:
        return None
    mid = (first + last) // 2
    # 해당 원소를 가지면서 가장 왼쪽에 있는 경우 index 반환
    if (mid == 0 or arr[mid - 1] < target) and target == arr[mid]:
        return mid
    elif arr[mid] < target:
        return findFirstIndex(arr, target, mid + 1, last)
    elif arr[mid] >= target:
        return findFirstIndex(arr, target, first, mid - 1)


def findLastIndex(arr, target, first, last):
    if first > last:
        return None
    mid = (first + last) // 2
    # 해당 원소이면서 가장 오른쪽에 있는 경우 index 반환
    if (mid == n - 1 or arr[mid + 1] > target) and target == arr[mid]:
        return mid
    elif arr[mid] <= target:
        return findLastIndex(arr, target, mid + 1, last)
    else:
        return findLastIndex(arr, target, first, mid - 1)


n, x = map(int, input().split())
array = list(map(int, input().split()))
count = countValue(array, x)
if count == 0:
    print(-1)
else:
    print(count)

 

2. 라이브러리 bisect를 사용한 정답 코드

import bisect

def countValue(arr, leftValue, rightValue):
    left_idx = bisect.bisect_left(arr, leftValue)
    right_idx = bisect.bisect_right(arr, rightValue)

    return right_idx - left_idx


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

count = countValue(array, x,x)

if count == 0:
    print(-1)
else:
    print(count)

 

bisect

파이썬에서 이진탐색을 쉽게 구현하기 위해 제공하는 라이브러리

 

'정렬된 배열'에서 특정 범위에 속하는 원소를 찾아야 할 때 효과적으로 사용된다.

bisect 라이브러리의 bisect_left(), bisect_right() 함수를 이용한다.

 

 

예를 들어, 정렬된 리스트 a = [1,5,5,5,8]가 존재한다고 가정하자.

이때 bisect_left(a, 5)와 bisect_right(a, 5)는 각각 인덱스 값으로 1, 4를 반환한다.