문제

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 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. 여태까지 지나온 스테이지들 중에서, 적이 가장 많은 스테이지들에만 무적권을 사용하고 싶다.
2. 무적권을 사용하면 병사의 수가 차감되지 않는다. (k만 차감)

여기에서 1번 아이디어가 되게 핵심인데, 즉, 우리는 지나온 스테이지들 중에서 최댓값만 한번에 쏙쏙 골라서 무적권을 적용시키고 싶다는 것이다. 그렇게 하려면, 최대힙을 사용하면 된다!!!

다음은 알고리즘이다.
enemy 리스트를 순회하면서 다음과 같은 로직을 수행한다.

• 가능하면 바로 병사를 소모한다. 이때 적의 수를 최대힙에 넣어준다.
• 만약, 바로 병사를 소모할 수 없는 상황이라면 아래 내용대로 행동한다.
  • [k > 0]
    • 현재 적의 수를 최대힙에 넣는다.
    • 현재 n에다가 heappop을 해서 나온 수를 더한다. (즉, 여태까지 지나온 스테이지들에서 가장 적이 많이 나온 구간에 대한 보상을 해줌으로써 무적권을 사용한다는 의미)
    • 따라서 k를 1 감소시켜준다.
  • [k == 0]
    • 희망이 없다. 게임 종료. (반복문 break)
  • answer += 1

코드

from heapq import heappush, heappop
def solution(n, k, enemy):
    rounds = len(enemy)
    #무적권이 충분한 경우
    if k >= rounds:
        return rounds
    #그렇지 않은 경우
    answer = 0
    heap = []
    
    for ene in enemy:
        #일단 병사 소모할 수 있으면 소모하여 heappush(최대힙이니까 음수로 저장)
        if n >= ene:
            n -= ene
            heappush(heap, -ene)
        
        #병사 소모할 수 없다면 (n < ene)
        else:
            #근데 무적권이 있다면, 무적권을 소모하고 n에 heap[0]을 더한 후, 현재 ene를 막는 데 쓴다.
            #그리고 heappush(heap, -ene) 한다.
            if k > 0 and heap:
                heappush(heap, -ene)
                n -= heappop(heap)
                k -= 1
                n -= ene
            #무적권은 있는데 여태까지 heap에 쌓아둔 게 없다면, 무적권만 소모한다.
            elif k > 0 and not heap:
                k -= 1
            #병사도 없고 무적권도 없다면, 겜 끝
            else:
                break
        
        answer += 1
    
    return answer