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

Problem : 코딩테스트 연습 - 디펜스 게임 | 프로그래머스 스쿨 (programmers.co.kr)

Status : Solved

Time : 00:16:01

 

---

 

문제 설명

 

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

 

 

---

 

풀이

잘 생각해보면, 현재까지의 stage중 가장 enemy가 많았던 k개에 무적권을 쓰는 것이 항상 가장 최선의 전략임을 알 수 있다. 따라서 그리디하게 가장 많은 enemy를 저장하며 문제를 풀이하면 된다. 이 경우 priority queue, heap를 쓰는 게 좋다.

즉 주요 알고리즘은 다음과 같이 나타낼 수 있다.

1. stage 개수보다 무적권의 개수 k가 같거나 더 많다면 모든 경우에 무적권을 사용할 수 있으므로, stage의 길이를 반환한다.

2. priority queue에 초반 k의 enemy를 저장한다.

3. k를 넘어선다면, i번째 stage의 enemy[i-1]를 priority queue에 저장하고, 하나를 pop한다. 그리고 그 pop한 enemy를 남은 병사 수 n에서 차감한다. 만약 차감하는 데 실패한다면(즉 차가 0보다 작아진다면) 최선의 경우에도 더 이상 게임을 진행하는 것이 불가능하므로, i-1를 리턴한다.

 

풀이 코드

from heapq import heappush, heappop

def solution(n, k, enemy):
    stage = len(enemy)
    if k >= stage :
        return stage
    q = []
    
    for i in range(stage) :
        heappush(q, enemy[i])
        if len(q) > k :
            last = heappop(q)
            if last > n :
                return i
            n -= last
        
    return stage

풀이 완료!