[이진 탐색] 이코테 (파이썬) 정렬된 배열에서 특정 수의 개수 구하기 풀이

[문제]

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있습니다. 이때 이 수열에서 x가 등장하는 횟수를 계산하세요.

 

예를 들어 수열 {1, 1, 2, 2, 2, 2, 3}이 있을 때 x = 2라면, 현재 수열에서 값이 2인 원소가 4개이므로 4를 출력합니다.

 

단, 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 시간 초과 판정을 받습니다.

 

[입력 조건]

1. 첫째 줄에 N과 x가 정수 형태로 공백으로 구분되어 입력됩니다. (1 <= N <= 1,000,000)

 

2. 둘째 줄에 N개의 원소가 정수 형태로 공백으로 구분되어 입력됩니다.

 

[출력 조건]

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

<입력 예시>
7 2
1 1 2 2 2 2 3
<출력 예시>
4

 

[풀이]

from bisect import bisect_left, bisect_right

def count(arr, left_value, right_value) :
  right_index = bisect_right(arr, right_value)
  left_index = bisect_left(arr, left_value)
  return right_index - left_index

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

result = count(arr, x, x)

if result != 0 :
  print(result)
else :
  print(-1)

 

이 문제는 파이썬 표준 라이브러리인 bisect_right( )와 bisect_left( )를 사용하여 쉽게 해결할 수 있다.

 

먼저 데이터의 개수(n)와 찾고자 하는 값(x)을 입력받고 전체 데이터를 arr에 입력한다.

 

count( ) 메소드에 데이터들을 전달하여 수행한 뒤 반환받은 값을 변수 result에 할당하여 

 

만약 그 값이 0이 아니라면 그대로 출력하고 0이라면 -1을 출력한다.

 

count( ) 메소드에서는 변수 right_index에 bisect_right( )를 사용하여 해당 인덱스를 할당하고

 

변수 left_index에 bisect_left( )를 사용하여 해당 인덱스를 할당한 뒤 right_index에서 left_index를 뺀 값을 반환한다.

 

---

출처

이것이 코딩 테스트다 with 파이썬 - 나동빈 저