디펜스 게임

from queue import PriorityQueue def solution(n, k, enemy): # cumsum: n번째 라운드까지의 누적합 cumsum, c = [], 0 for x in enemy: c += x cumsum.append(c) # ksum: n번째 라운드까지 따졌을 때 가장 큰 k개의 합 ksum, c = [], 0 q = PriorityQueue() # min-heap for x in enemy: if q.qsize() < k: q.put(x) c += x ksum.append(c) else: if q.queue[0] < x: # 지금의 enemy 수가 더 많음 c = c - q.get() + x q.put(x) ksum.append(c) remain = [n - (a-b) for a, b in zip(cumsum, ksum)] answer = 0 while answer < len(enemy) and remain[answer] >= 0: answer += 1 return answer