[Python, Java] 프로그래머스 - 디펜스게임

문제

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

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

준호는 무적권을 적절한 시기에 사용하여 최대한 많은 라운드를 진행하고 싶습니다.

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

 

n : 처치 가능한 적의 수, k : n 소모 없이 넘기기 가 주어졌을 때 최대로 살아남을 수 있는 라운드를 구하는 문제입니다.

풀이

1. n을 소모해가며 enemy를 순회합니다.
2. 현재의 enemy가 n보다 커져을때, 현재까지 만난 enemy 중 가장 큰 enemy를 대상으로 k를 소모합니다. 즉, k-= 1을 수행하고, n에 가장 큰 enemy(단, 현재까지 순회한 enemy 대상)만큼 더하면 됩니다.
3. 더이상 k가 없거나, enemy를 전부 순회했으면 종료됩니다.

처음엔 정렬과정을 생각하지 못하고 단순히 max를 계산하여 더하였습니다. 당연하게도 이는 시간초과가 나는 테스트 케이스들이 발생합니다.

def solution(n, k, enemy):
    answer = 0
    history = []
    for e in enemy:
        n -= e
        history.append(e)
        
        if n < 0:
            if k > 0:
                largest_e = max(history)
                n += largest_e
                history.remove(largest_e)
                k -= 1
            else:
                break
        answer += 1
                
    return answer

어떻게 하면 정렬한 상태로 가장 큰 값을 뽑아낼 수 있을까 생각하다가 리스트가 생각이 났고, 최대힙을 사용해야겠다는 결론을 내렸습니다. 최대힙을 사용해서 같은 논리로 문제를 해결할 수 있습니다.

1. 현재까지 만난 적을 저장할 리스트를 생성합니다. (최대힙으로 사용)
2. enemy를 순회하며 현재 enemy만큼 n을 감소시킵니다.
3. 리스트에 현재만난 enemy를 집어넣는데, 최대힙을 사용합니다.
4. 만약 k 가 남아 있고, n이 0보다 작아졌다면 k를 감소시키고 최대힙에서 가장 큰 값을 제거합니다. 제거한 값만큼 n에 값을 더합니다.
5. k가 남아있지 않다면 종료하고, 아니라면 정답 카운터에 1을 더합니다.

전반적인 논리는 앞서 작성했던 방법과 같습니다. 차이라면 max를 사용했을 때와 최대힙을 사용했을 때에 있을 것 같아요. max는 O(N)을 최대힙의 push와 pop모두 O(logN) 입니다. 자바로 풀이할 경우 우선순위 큐를 사용하여 같은 로직으로 해결할 수 있습니다.

 

Python

import heapq

def solution(n, k, enemy):
    answer = 0
    max_heap = []

    for e in enemy:
        n -= e
        heapq.heappush(max_heap, -e)
        if n < 0:
            if k > 0:
                n += -heapq.heappop(max_heap)
                k -= 1
            else:
                break
        answer += 1

    return answer

Java

import java.util.PriorityQueue;

class Solution {
    public int solution(int n, int k, int[] enemy) {
        int answer = 0;
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);

        for (int e : enemy) {
            n -= e;      
            maxHeap.xadd(e);

            if (n < 0) {
                if (k > 0) {
                    n += maxHeap.poll();
                    k -= 1;
                } else {
                    break;
                }
            }
            answer += 1;
        }
        return answer;
    }
}