[백준] 10816번 숫자카드2, 시간초과 - 이분탐색만으로도 안됨, bisect 모듈!

https://www.acmicpc.net/problem/10816

10816번: 숫자 카드 2

 

이분 탐색으로 아래 풀이 처럼 풀었는데도 시간초과가 뜬다.

# BOJ_10816
# 입력
n = int(input())
cards = list(map(int, input().split()))
m = int(input())
show_cards = list(map(int, input().split()))


# 두 카드 목록 모두 오름차순으로 정렬한다.
cards.sort()


# 카드 갯수를 세는 함수
def count_card(cards, show_card):
    # 이분탐색으로 제시된 카드 찾는다.
    start, end = 0, n - 1
    while start <= end:
        mid = (start + end) // 2
        if cards[mid] == show_card:
            # show_card가 시작되는 부분, 끝나는 부분 인덱스를 각각 구한다.
            card_start, card_end = mid, mid
            while cards[card_start] == show_card:
                if card_start > 0:
                    card_start -= 1
                else:
                    card_start -= 1
                    break
            while cards[card_end] == show_card:
                if card_end < len(cards) - 1:
                    card_end += 1
                else:
                    card_end += 1
                    break
            return card_end - card_start - 1
        elif cards[mid] > show_card:
            end = mid - 1
        else:
            start = mid + 1
    return 0


# 카드 갯수를 각각 출력해준다.
for show_card in show_cards:
    print(count_card(cards, show_card))

 

---

 

원인을 찾던 중에 upper bound / lower bound 개념을 추가해야 한다는 걸 알았다.

https://velog.io/@himinhee/%EC%95%8C%EA%B3%A0%EB%A6%AC%EC%A6%98-Lower-bound-Upper-bound-python

Lower bound, Upper bound (python)

그리고 위 개념을 한번에 쓸 수 있는 모듈: bisect 모듈

https://velog.io/@himinhee/python-bisect%EB%AA%A8%EB%93%88

[python] bisect모듈

 

 

bisect 모듈을 이용하여 아래와 같이 구현하니 맞았다.

 

기존 

if cards[mid] == show_card:
            # show_card가 시작되는 부분, 끝나는 부분 인덱스를 각각 구한다.
            card_start, card_end = mid, mid
            while cards[card_start] == show_card:
                if card_start > 0:
                    card_start -= 1
                else:
                    card_start -= 1
                    break
            while cards[card_end] == show_card:
                if card_end < len(cards) - 1:
                    card_end += 1
                else:
                    card_end += 1
                    break
            return card_end - card_start - 1

 

bisect으로 시간 복잡도 줄인 코드

if cards[mid] == show_card:
    # show_card가 시작되는 부분, 끝나는 부분 인덱스를 각각 구한다.
    card_start = (bisect.bisect_left(cards, show_card))
    card_end = (bisect.bisect_right(cards, show_card))
    return card_end - card_start

두 코드 모두 같은 기능(특정 숫자의 시작과 끝 범위 찾기)을 구현한 코드이다.

 

 

# BOJ_10816
import bisect

# 입력
n = int(input())
cards = list(map(int, input().split()))
m = int(input())
show_cards = list(map(int, input().split()))


# 두 카드 목록 모두 오름차순으로 정렬한다.
cards.sort()


# 카드 갯수를 세는 함수
def count_card(cards, show_card):
    # 이분탐색으로 제시된 카드 찾는다.
    start, end = 0, n - 1
    while start <= end:
        mid = (start + end) // 2
        if cards[mid] == show_card:
            # show_card가 시작되는 부분, 끝나는 부분 인덱스를 각각 구한다.
            card_start = (bisect.bisect_left(cards, show_card))
            card_end = (bisect.bisect_right(cards, show_card))
            return card_end - card_start
        elif cards[mid] > show_card:
            end = mid - 1
        else:
            start = mid + 1
    return 0


# 카드 갯수를 각각 출력해준다.
for show_card in show_cards:
    print(count_card(cards, show_card))