[Programmers] 디펜스 게임

디펜스 게임

문제 설명

준호는 무적권을 적절한 시기에 사용하여 최대한 많은 라운드를 진행하고 싶습니다.
준호가 처음 가지고 있는 병사의 수 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 해주세요.

풀이

나는 힙을 생각하지 못해서 풀지 못했다. enemy의 길이가 최대 100만이므로 시간복잡도 O(NlogN)까지 가능하다. 최대힙 또는 최소힙으로 풀 수 있다.

1. 최소힙에 enemy를 순서대로 넣는다.
2. 최소힙의 크기가 k보다 크면 가장 작은 값을 꺼내고 sum에 더한다.
3. sum>n 이 되면 적의 인원이 더 많아졌다는 것이고 게임이 종료된다.
4. 게임이 중간에 종료되지 않았다면 모든 라운드를 통과한 것이다.

코드

#include <bits/stdc++.h>
using namespace std;

int solution(int n, int k, vector<int> enemy) {
    int sum=0;
    priority_queue<int, vector<int>, greater<int>> pq; // 최소힙
    
    for(int i=0; i<enemy.size(); i++){
        pq.push(enemy[i]);
        
        if(pq.size() > k){
            sum+=pq.top();
            pq.pop();
        }
        if(sum > n) // 게임 종료
            return i;
    }
    return enemy.size(); // 모든 라운드 통과
}

정리

참고