[프로그래머스] 디펜스 게임 Lv.2 Python

문제 링크

1. 문제 설명

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 n명으로 연속되는 적의 공격을 순서대로 막는 게임입니다. 디펜스 게임은 다음과 같은 규칙으로 진행됩니다.

준호는 처음에 병사 n명을 가지고 있습니다.
매 라운드마다 enemy[i]마리의 적이 등장합니다.
남은 병사 중 enemy[i]명 만큼 소모하여 enemy[i]마리의 적을 막을 수 있습니다.
예를 들어 남은 병사가 7명이고, 적의 수가 2마리인 경우, 현재 라운드를 막으면 7 - 2 = 5명의 병사가 남습니다.
남은 병사의 수보다 현재 라운드의 적의 수가 더 많으면 게임이 종료됩니다.
게임에는 무적권이라는 스킬이 있으며, 무적권을 사용하면 병사의 소모없이 한 라운드의 공격을 막을 수 있습니다.
무적권은 최대 k번 사용할 수 있습니다.
준호는 무적권을 적절한 시기에 사용하여 최대한 많은 라운드를 진행하고 싶습니다.

준호가 처음 가지고 있는 병사의 수 n, 사용 가능한 무적권의 횟수 k, 매 라운드마다 공격해오는 적의 수가 순서대로 담긴 정수 배열 enemy가 매개변수로 주어집니다. 준호가 몇 라운드까지 막을 수 있는지 return 하도록 solution 함수를 완성해주세요.

제한사항

• 1 ≤ n ≤ 1,000,000,000
• 1 ≤ k ≤ 500,000
• 1 ≤ enemy의 길이 ≤ 1,000,000
• 1 ≤ enemy[i] ≤ 1,000,000
  enemy[i]에는 i + 1 라운드에서 공격해오는 적의 수가 담겨있습니다.
• 모든 라운드를 막을 수 있는 경우에는 enemy[i]의 길이를 return 해주세요.

• 입출력 예

n	k	enemy	result
7	3	[4, 2, 4, 5, 3, 3, 1]	5
2	4	[3, 3, 3, 3]	4

2. 문제 접근

무적권을 이용해 최대 k번 공격을 넘길 수 있지만 대부분의 경우 k가 주어진 배열 enemy의 길이보다 크지 않습니다. 따라서 k가 enemy 배열의 길이보다 작아서 무적권을 선택적으로 사용해야 하는 경우를 고려해야 합니다.
현재 라운드가 $k + 1$ 번째 라운드라 합시다. $k + 1$ 번의 enemy 웨이브 중에 가장 어려운 $k$개의 무적권을 사용한다고 가정해도 가장 약한 한 번의 라운드는 병사를 소모할 수 밖에 없습니다. 따라서 $k + 1$ 번째부터 이후 계속 가장 약한 라운드마다 병사를 소모한다고 가정하며 진행해야 합니다. 이럴 때 사용할 수 있는 가장 좋은 자료구조는 우선순위배열, 힙(heap)입니다.

3. 문제 풀이

from heapq import heappush, heappop

def solution(n, k, enemy):
    hq = []
    for round, monster in enumerate(enemy):
        heappush(hq, monster)
        if len(hq) > k:
            n -= heappop(hq)
        if n < 0:
            return round

    return len(enemy)

우선 enemy 배열의 원소를 heap에 하나씩 넣어 줍니다. 만약에 k가 enemy보다 크다면(if len(hq) < k:) 바로 enemy 배열의 길이를 return해주면 됩니다. heap의 길이가 k보다 커질 경우 매번 heap에서 가장 작은 원소를 제거하고 n(병사의 수)에서 빼줍니다. n이 음수가 된다면 더 이상의 라운드를 진행할 수 없는 상태가 됩니다. 따라서 그 전 라운드의 숫자를 반환해 주면 되는데, enumerate를 통해 인덱스를 호출하면 라운드의 순서를 찾을 수 있습니다.