[프로그래머스] 디펜스 게임 (Python)

Source

https://school.programmers.co.kr/learn/courses/30/lessons/142085

프로그래머스

 

Code

import heapq

def solution(n, k, enemy):
    heap = []
    answer = 0
    sum = 0

    for i in enemy:
        heapq.heappush(heap, (-i, i))
        sum += i
        answer += 1

        if sum > n:
            if k:
                sum -= heapq.heappop(heap)[1]
                k -= 1
            else:
                return answer - 1

    return answer

 

How to solve?

1. heapq를 선언하고 (-i, i) 형태의 튜플을 삽입하는 것으로 최대힙을 구현한다.

 

2. 매번 배열의 값을 더하는 것 보다 하나의 sum으로 합계를 구하는 것이 더 빠르기 때문에

sum = 0을 선언한 이후 n보다 클 경우 무적권을 사용 해준다.

 

3. sum > n이고, 만약 k가 남아있을 경우 heapq.heappop을 통해 최대 값을 제거 해준다.

 

4. sum > n이고, 만약 k가 남아있지 않을 경우 answer에 -1을 해준 값을 리턴한다.

(sum > n이기 때문에 막아내지 못하고, 무적권이 없기 때문에 그 이전 라운드가 최대값이다)

 

5. 모든 for문을 순회 한 경우 answer을 리턴 해준다.