[백준 - 1202] 보석 도둑 - C++

분류: 그리디, 자료구조, 시간 복잡도 /

 

문제

문제 링크

문제 설명

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 N개 있다. 각 보석은 무게 M_i와 가격 V_i를 가지고 있다. 상덕이는 가방을 K개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 C_i이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000)

다음 N개 줄에는 각 보석의 정보 M_i와 V_i가 주어진다. (0 ≤ M_i, V_i ≤ 1,000,000)

다음 K개 줄에는 가방에 담을 수 있는 최대 무게 C_i가 주어진다. (1 ≤ C_i ≤ 100,000,000)

모든 숫자는 양의 정수이다.

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

 

풀이

#include <iostream>
#include <queue>
#include <algorithm>
#define M first
#define V second
using namespace std;

int N, K;
pair<int,int> jewels[300000]; 
int capacity[300000];

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    cin >> N >> K;
    for(int i=0; i<N; i++) cin >> jewels[i].M >> jewels[i].V;
    for(int i=0; i<K; i++) cin >> capacity[i];
    sort(jewels, jewels+N);
    sort(capacity, capacity+K);
    
    long long ans = 0;
    priority_queue<int, vector<int>, less<int>> pq;
    for(int bag=0, i=0; bag<K; bag++) {
        while(i < N and jewels[i].M <= capacity[bag]) 
            pq.push(jewels[i++].V);
        
        if(pq.empty()) continue;
        ans += pq.top();
        pq.pop();
    }
    cout << ans;
}

 

용량이 작은 가방부터 현재 가방에 넣을 수 있는 가장 비싼 보석을 담는다.

 

• 보석을 무게 오름차순으로, 가방을 용량 오름차순으로 정렬한다.
• 용량이 작은 가방부터 현재 가방에 넣을 수 있는 모든 보석을 `가격 max heap`에 넣는다.
  • 넣을 수 있는 보석이 없으면 다음 가방에 대해 윗 줄을 수행. (`pq.empty()`)
  • 있다면 현재 가방에 넣을 수 있는 가장 비싼 보석을 가방에 넣는다. (`pq.top()`)

`pq`에는 현재 가방에 넣을 수 있는 보석 중 앞에서 다른 가방에 넣지 않은 보석들이 누적해서 남아있어 보석도 선형으로 탐색 가능하다.

 

`N, K`가 모두 최대 300000으로 `O(N * K)`의 풀이는 시간초과가 발생한다.

여기서 log가 포함된 풀이를 해야한다는 것을 알 수 있고 정렬 혹은 비선형 자료구조를 사용해야한 다는 것을 알 수 있다.

 

위 풀이의 시간 복잡도를 보면 정렬에서 `O(NlogN + KlogK)`,

그리디 for문에서 `O(K + NlogN)`으로 `O(NlogN + KlogK)`의 시간복잡도를 갖는다.

그리디 for문은 for문 안에 while문이 있어 `O(KNlogN)`처럼 보이지만 `bag`과 `i`가 독립적으로 증가하기 때문에 `O(K + NlogN)`이다.

 

- 이해를 위한 시간초과 예시

bool used[300000];

long long ans = 0;
for(int bag=0; bag<K; bag++) {
    int maxIdx = -1;
    int maxVal = 0;
    for(int i=0; i<N; i++) {
        if(used[i]) continue;
        if(jewels[i].M > capacity[bag]) break;
        if(jewels[i].V <= maxVal) continue;
        maxVal = jewels[i].V;
        maxIdx = i;
    }
    if(maxIdx == -1) continue;
    used[maxIdx] = true;
    ans += maxVal;
}
cout << ans;

max heap을 사용하지 않고 매번 for문을 돌며 현재 가방에 담을 수 있는 보석중 아직 가방에 담지 않은 보석에서 가치가 가장 높은 보석을 찾는다면 `O(NK)`로 시간초과다.