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

Lv2, Greedy, heapq

문제 설명

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 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 해주세요.

입출력 예시

풀이

1. 오답 풀이
   -> deque를 통해 내림차순으로 정렬 한 후, 해결하고자 했다.
   TC를 6개나 만들어서 해봤는데 TC는 전부 성공했었는데 제출을 하면 46.9 / 100.0 으로 실패했다.

from collections import deque
def solution(n, k, enemy):
    answer = 0
    enemy.sort(reverse=True)
    q = deque(enemy)
    
    while 1:
        while k > 0:
            q.popleft()
            k -= 1
            answer += 1
            if len(q) == 0:
                return answer
        
        if len(q) > 0:
            if q[0] <= n:
                n -= q[0]
                q.popleft()
                answer += 1
            elif q[0] > n:
                break
        else:
            break

    return answer

2. 정렬을 하면 안되는 것 같아서 heapq를 써보기로 했다.

import heapq

def solution(n, k, enemy):    
    if k >= len(enemy):
        return len(enemy)
    
    q = []
    
    for i in range(len(enemy)):
        heapq.heappush(q, enemy[i])     
        # 힙 내부의 사람 수 만큼 k를 쓰면 됨
        if len(q) > k:
            last = heapq.heappop(q)
            if last > n:
                return i        # 더이상 라운드 진행 불가
            n -= last

    return len(enemy)

문제 자체는 쉬웠던거 같은데 처음에 접근을 잘못하다보니 다른 풀이 방법이 생각나지 않았다.