보석 도둑

문제: 보석 도둑 (백준 1202번)

구조 분석

구조는 전형적인 greedy 알고리즘이었다. 일반적인 greedy fractional knapsack 문제에서 가방이 여러개 있는 문제여서 가장 작은 가방에 가장 비싼 보석부터 넣는 방식으로 알고리즘을 작성하면 된다.

코드

#include <stdio.h>
#include <iostream>
#include <queue>
#include <algorithm>
#define MAX 300000
using namespace std;

int n, k;
int C[MAX];
pair<int,int> J[MAX];
priority_queue<int> pq;

int main() {
    cin >> n >> k;
    for (int i=0; i<n; i++) {
        cin >> J[i].first >> J[i].second;
    }
    for (int i=0; i<k; i++) {
        cin >> C[i];
    }

    sort(J,J+n);
    sort(C,C+k);

    long long cost = 0;
    int idx = 0;
    for (int i=0; i<k; i++) {
        while (J[idx].first <= C[i] && idx < n) {
            pq.push(J[idx].second);
            idx++;
        }
        if (!pq.empty()) {
            cost += pq.top();
            pq.pop();
        }
    }
    cout << cost << "\n";
    return 0;
}

풀이 방법

처음에는 그냥 vector를 이용해서 이중 for문으로 돌렸다. 그랬더니 시간복잡도가 O(N*K)까지 늘어나면서 시간초과가 되었다. 그래서 이를 해결하기 위한 방법을 찾아보니 priority queue를 사용해야 했다.
priority queue를 사용해서 아래의 방식으로 구현하면 된다.

1. 보석이랑 가방을 오름차순으로 정렬한다. (무게가 낮은 순서)
2. 가장 작은 가방부터 보석이랑 무게를 비교하고, 보석이 들어가지면 priority queue에 push를 실시한다.
3. push가 모두 끝났으면 queue의 top을 cost에 더하고 pop을 실시한다.
4. 위의 과정을 반복한다.

정리

시간초과 때문에 좀 당황한 문제이다. priority queue를 사용하여 이를 O(N+K)에 탐색을 끝낼 수 있다. priority queue를 사용하는 데에 익숙해져야겠다.