[heap/백준] 1417번 : 국회의원 선거

- 문제 : https://www.acmicpc.net/problem/1417

 

문제 이해 

- 후보 1번을 포함한 n명의 후보의 예상 투표 개수가 주어지는데, 1번 후보가 가장 큰 값을 가질 수 있게 만드는 문제다

- 1번이 가장 많은 예상 득표수를 가져야하므로, 2번 - n번의 후보들 중 가장 득표수가 많은 후보를 뽑아 그 후보의 득표수를 줄이면 된다

- heap을 사용하는 이유는 2번 ~ N번 후보에서 가장 높은 득표수 후보를 계속 뽑아야 하므로 가장 작은/큰 값이 가장 앞단에 위치하고, 시간 복잡도가 O(logN)으로 짧은 자료구조가 필요했기 때문이다. 

- 이때 최대 힙을 구현해야하므로 득표수는 마이너스 값으로 진행한다 (가장 작은 마이너스 값이 가장 큰 양수값이므로)

- 다른 후보에서 뺏은 득표수를 1번 후보 득표수에 계속 더하고, 1번 후보의 득표수가 다른 어떤 후보 보다 더 클 때 멈추면 된다

- 마지막으로 득표수를 뺏는 반복을 몇 번 했는지 세서 프린트 하면 끝

 

풀이 로직

- 1번 후보는 계속 써야하므로 first 변수에 저장하고, votes에는 나머지 후보들의 마이너스 득표 수를 추가한다

- votes를 heap으로 진행할 것이므로 heapify도 해준다

- while문을 돌면서 votes에서 heappop을 해서 후보 중 가장 큰 득표수를 가져온다

- 만약 first의 값이 더 크다면 first 값이 최대이므로 반복문을 break하고

- 아니라면 해당 후보의 득표수를 줄이고(vote_num + 1 (vote_num이 음수이기 때문), 1번 후보의 득표수를 늘린다 (first -= 1)

 

 

정답 풀이

import heapq

n = int(input())

first = -1 * int(input())
votes = []
for i in range(1, n):
    votes.append([-1 * int(input()), i])

heapq.heapify(votes)

count = 0
while votes:
    vote_num, number = heapq.heappop(votes)
    
    if vote_num > first: # 1번 득표수가 제일 클 때 
        break
        
    heapq.heappush(votes, [vote_num + 1, number])
    first -= 1
    count += 1
    
print(count)