[프로그래머스 Level2] 디펜스 게임 - 파이썬(Python)

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

프로그래머스

문제 설명

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

보다시피 시간 초과에 조심해야하는 문제이다..!

 

우선,  k개의 무적권은 될 수 있는 만큼 모두 사용하는 것이, 특히 여러 라운드 내에 적의 숫자가 가장 큰 숫자에 사용하는 것이 좋다고 생각이 들었다.

따라서 minHeap을 두고 처음에 k개의 요소를 담고 시작하고, 각각 다음 라운드의 숫자를 넣은 후 힙의 가장 앞에 있는(그 중 가장 작은) 값을 빼서 n에 빼주는 방식으로 진행했다.

 

 

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

위 예제에서의 과정)

i=0,1,2,: damage= 0... k를 모두 사용

i=3: damage= 2... k로 4, 4, 5를 막음

i=4: damage= 2+3= 5... k로 4, 4, 5를 막음

i=5: damgae= 2+3+3= 8   *** n보다 큰 수, return i ***

 

 

<< 전체 코드 >>

import heapq

def solution(n, k, enemy):
    minHeap = enemy[:k]
    heapq.heapify(minHeap)
    for i in range(k, len(enemy)):
        heapq.heappush(minHeap, enemy[i])
        n -= heapq.heappop(minHeap)
        if n < 0:
            return i
    return len(enemy)