[BOJ] 1417: 국회의원 선거 (Python)

문제

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

풀이 전략

Priority Queue

우선순위 큐를 이용한 간단한 문제다.

다솜이의 표를 따로 변수를 지정해 관리하고(dasom), 나머지 표는 list 자료구조에 집어넣는다. (heap)

다음과 같은 순서로 문제를 푼다.

heapq 메서드를 이용해 리스트를 우선순위 큐로 바꾼 후, 마개조를 통해 (파이썬에서는 heapq가 최소 힙을 만드므로, 튜플을 만들어 첫 번째 요소는 값에 -를 붙이고, 두 번째 요소는 원래 값을 집어넣어 정렬은 첫 번째 값을 기준으로 하되, 실제로는 두 번째 값을 이용하도록 한다) 최대 힙을 구현한다.

힙 안의 요소를 하나씩 꺼내보며 다솜이의 표보다 크거나 같을 경우 -1을 해 다시 집어넣는다 + count 변수를 지정해 이 연산을 할 때마다 1씩 더해준다
다솜이의 표에 +1을 하는 것을 반복한다.

만약 다솜이의 표가 힙 안의 최댓값보다 크다면, 종료하고 count를 출력한다.

우선순위 큐에서 삽입, 삭제 연산은 각각 $O(\log n)$ 의 시간복잡도를 지니고, 최대 n-1번의 연산을 해야 하므로, 총 시간복잡도는 $O(2(n-1) \log n) = O(n \log n)$ 이다.

구현

import heapq

if __name__ == "__main__":
    n = int(input())
    dasom = int(input())

    heap = []
    for _ in range(n - 1):
        num = int(input())
        heapq.heappush(heap, (-num, num))


    if len(heap) == 0:
        print(0)

    else:
        count = 0
        while True:
            now = heapq.heappop(heap)[1]
            if now >= dasom:
                count += 1
                dasom += 1
                heapq.heappush(heap, (-now + 1, now - 1))
            else:
                break

        print(count)