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

문제

문제 설명

• n, k, 숫자 리스트가 주어진다.
• 리스트를 순회하며 값을 누적하는데 누적 연산은 최대 k번 건너뛸 수 있다.
• 누적값이 n을 넘지 않도록 하여야 한다.
• 누적 연산이 마무리 된 인덱스 + 1값을 반환하여야 한다.

풀이

• 리스트의 값을 차례대로 순회하며 최대 힙에 저장 & 누적 값을 연산한다.
• 누적 값이 n이 넘어간다면, k를 고려하며 힙의 최대값을 제거한다.

import heapq

def solution(n, k, enemy):
    
    heap = []
    result = 0
    cnt = 0
    for e in enemy:
        heapq.heappush(heap, -e)
        cnt += e
        if cnt > n:
            if not k:
                return result
            cnt += heapq.heappop(heap)
            k-=1
        result += 1
                
    return result