[프로그래머스] 디펜스 게임 - JAVA

문제 링크

문제 정의

• n명의 병사
• k개의 무적권
• 무적권을 소모하면 해당 라운드의 적들을 병사의 소모없이 제거
• 라운드 별 적의 수를 담은 배열 enemy
• n명의 병사와 k개의 무적권을 사용해 매 라운드 enemy만큼 오는 적들을 최대한 막을 수 있는 라운드의 수를 구하라

시간 복잡도

이 문제의 입력 제한사항들은 아래와 같습니다.

• $1 ≤ n ≤ 1,000,000,000$
• $1 ≤ k ≤ 500,000$
• $1 ≤ enemy의 길이 ≤ 1,000,000$
• $1 ≤ enemy[i] ≤ 1,000,000$

enemy의 길이가 $1,000,000$이므로 최대한 순회하면서 문제를 풀 수 있어야 합니다.
라운드 수를 $N$이고 무적권의 수를 $K$라고 가정했을 때
최대로 가능한 시간 복잡도는 $O(KlogN)$이라고 생각했습니다.

접근 방법

문제를 보자마자 최대한 라운드를 진행하기 위해서는 다음과 같은 사고가 필요하다고 생각했습니다.

1. 일단 무적권을 쓰지말고 최대한 라운드를 진행해보자
2. 더 이상 라운드를 진행할 수 없을 때 무적권을 사용하자
3. 무적권을 사용할 때는 진행한 라운드 내에서 가장 적이 많이나온 라운드에 사용하자
4. 무적권이 없고 더 이상 라운드를 진행할 수 없을 때 정답을 도출하자

위와 같은 생각을 가장 효율적으로 사용할 수 있는 자료구조는 우선순위 큐라고 생각했습니다.
진행한 모든 라운드의 적의 수를 우선순위 큐 pq에 저장하고
무적권을 사용할 때는 내림차순으로 정렬된 pq에서 값을 뽑으면 가장 큰 수가 나오기 때문에
$O(logN)$ 시간 내에 문제를 해결할 수 있었습니다.

코드

코드는 간결하게 구현했습니다.
다른 알고리즘을 사용할 필요 없이 우선순위 큐만 사용한 코드입니다.

import java.util.*;

class Solution {
    public int solution(int n, int k, int[] enemy) {
        PriorityQueue<Integer> pq = new PriorityQueue<>((o1, o2) -> {
            return o2 - o1;
        });
        int L = enemy.length;
        int count = 0;
        
        for (int i = 0; i < L; i++) {
            count += enemy[i];
            pq.offer(enemy[i]);
            
            if (count > n) {
                while (k > 0 && count > n) {
                    k--;
                    count -= pq.poll();
                }
            }
            if (k == 0 && count > n) return i;
        }
        
        return L;
    }
}