binary search: 정렬된 배열에서 특정 수의 개수 구하기

문제정의

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

입력조건

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

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

입출력예시

# 입력
7 2
1 1 2 2 2 2 3
# 출력
4

내가짠코드

from bisect import bisect_left, bisect_right

n, x = map(int, input().split(" "))
lst = list(map(int, input().split(" ")))
lst = [1,1,2,2,2,2,3]
lst.sort(reverse=False)
x = 2
result = bisect_right(lst, x) - bisect_left(lst, x)

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

정답코드

from bisect import bisect_left, bisect_right

def count_by_range(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_by_range(arr, x, x)

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

고찰

  • 정답코드와 다를것 없이 잘짰다.
  • 파이썬에서 이렇게 유용한 함수를 잘 기억하자.

좋은 웹페이지 즐겨찾기