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

링크

디펜스 게임

문제의 조건을 확인해보면 문제가 원하는 답이 무엇인지는 추측할 수 있으나 접근하는 방법을 떠올리기 쉽지 않다.

문제에서 주어진 조건을 정리해보면 아래와 같다

• 준호는 초기 $n$ 명의 병사를 가진다.
• 매 라운드마다 enemy[i] 마리의 적이 등장한다
• 무적권은 $k$ 번 사용할 수 있다
• 행위
  • 현재 남은 병사보다 적이 더 많을 경우 게임은 종료
  • 현재 남은 병사보다 적이 더 적을 경우 게임을 진행
    • 무적권을 사용할 경우 병사의 수가 줄지 않는다
    • 무적권을 사용하지 않을 경우 병사의 수가 줄어든다.
• Target := 무적권을 최대 $k$ 번 사용하여 준호가 진행할 수 있는 최대 라운드는?

이제 Constraint도 고려해보자.

• 1 $\leq$ n $\leq$ $10^9$
• 1 $\leq$ k $\leq 5*10^5$
• $1 \leq$ len(enemy) $\leq 10^6$

enemy의 길이는 최대 $10^6$ 으로 $O(NlogN)$ 이하의 시간 복잡도로 문제를 해결 해야한다.

문제의 요구 사항을 수식으로 표현해보자.

• e: 적의 수k: 무적권 수
• n: 초기 병사 수

$k = 0$ 일 때

$e_1 + e_2 + e_3 + ... + e_m \leq n$ 이 되는 $m$의 최댓값

$k = 1$ 일 때

$e_1 + e_2 + e_3 + ... + e_m - (e_\alpha) \leq n$이 되는 $m$의 최댓값 (단, $\alpha$는 [1, m] 사이의 임의의 수)

여기서 $m$ 을 최댓값 이라고 가정한다면, $e_\alpha$는 $max(e_1, ..., e_m)$ 이라는 것을 어렵지 않게 알 수 있다.

이 부분은 $m$이 최댓값일 때, $e_a$ 가 최댓값이 아니라면 모순이 발생한다는 것을 통하여 보일 수 있다. 즉, $e_\alpha$ 가 max가 아니라면 $m < q$ 인 $q$에 대해서 $e_1 + ... + e_q$가 존재한다는 것을 보인다.

일반화하여 표현해보면 아래와 같이 표현된다.

$k = \beta$ 일 때

$e_1 + e_2 + ... + e_m - (e_{\delta_1} + e_{\delta_2} + ... + e_{\delta_k}) \leq n$ 이되는 $m$의 최댓값(단, $\delta_i$ 는 $(e_1, ..., e_m)$중 $i$번째로 가장 큰 수)

즉 이 문제의 핵심은 $(e_{\delta_1} + e_{\delta_2} + ... + e_{\delta_k})$를 빠르게 구하는 것에 있다.

$e_1, ...., e_n$까지 선형탐색하며, 빠르게 $(e_{\delta_1} + e_{\delta_2} + ... + e_{\delta_k})$을 최댓값으로 유지하는 방법은 min-heap 을 사용하여 $e_i$ 를 탐색할때마다 $e_i$ 와 min-heap의 top-element를 비교하여 최댓값으로 min-heap을 갱신하는 것이다.

그러므로, 전체 코드는 아래와 같다

#include <bits/stdc++.h>

using namespace std;

int solution(int n, int k, vector<int> enemy) {
    int answer = 0;

    long long sum = 0L;
    long long ac_max = 0L;

    priority_queue<int, vector<int>, greater<int>> PQ; // min heap
    int counter = 0;
    for (int i = 0; i < enemy.size(); i++) {
        int elem = enemy[i];
        sum += elem;

        if (PQ.size() < k) {
            ac_max += elem;
            PQ.push(elem);
        } else {
            int top = PQ.top(); PQ.pop();
            if (top < elem) {
                ac_max += (elem - top);
                PQ.push(elem);
            } else {
                PQ.push(top);
            }
        }

        if (sum - ac_max <= n) {
            counter++;
        } else {
            break;
        }
    }

    answer = counter;
    return answer;
}