[25.03.07] 보석 도둑

https://www.acmicpc.net/problem/1202

 

import java.util.*;
import java.lang.*;
import java.io.*;

// The main method must be in a class named "Main".
class Main {
    static int N;
    static int K;
    static PriorityQueue<int[]> qj;
    static PriorityQueue<Integer> bj;
    public static void main(String[] args) throws Exception{
        // N , K 받고
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] s= br.readLine().split(" ");
        // N 로 보석 무게 , 가격
        N = Integer.parseInt(s[0]);
        K= Integer.parseInt(s[1]);
        qj = new PriorityQueue<int[]>(new Comparator<int[]>(){
            @Override
            public int compare(int[] o1, int[] o2){
                if(o1[0]==o2[0]){
                    return o2[1]-o1[1];
                }
                return o1[0]-o2[0];
            }
        });
        bj = new PriorityQueue<Integer>();
        for(int i=0 ; i<N; i++){
            s = br.readLine().split(" ");
            int[] temp = new int[2];
            temp[0]=Integer.parseInt(s[0]);
            temp[1]=Integer.parseInt(s[1]);
            qj.add(temp);
        }

        // K 개 가방 무게
        for(int i=0; i<K; i++){
            int m = Integer.parseInt(br.readLine());
            bj.add(m);
        }


        // 가방 기준잡고
        long sum=0;
        PriorityQueue<Integer> price = new PriorityQueue<Integer>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2-o1;
            }
        });
        //가방 큐에 있으면
        while(!bj.isEmpty()){
            //가방 하나 가져와서
            int testBag = bj.poll();
            // 해당 가방 무게 >= 보석 무게 , 보석 큐에 보석 있을때
            while(!qj.isEmpty() && testBag>=qj.peek()[0]){
                price.add(qj.poll()[1]);
            }
            if(!price.isEmpty()){
                sum+=price.poll();
            }
            else{
                continue;
            }
        }

        System.out.printf("%d",sum);
        // 보석 하나씩 보면서 무게 들어가는지 안들어가는지
        // 들어간다면 push
        // 보석이 가방보다 커지면 우선순위에서 pop하고 sum
    }
}

 

와 상당히 해맸다. 

 

처음에는 

1) 가방 기준으로 for문 돌면서

2) 하나씩 보석을 기준잡은 가방에 넣어보면서 매번 그 가방에 넣을 수 있는 최대 가격의 보석을 넣으면 되는거 아닌가? >> 그리디

 

당연히 이러면 맞긴 하지만 가방과 보석의 이중 for문 때문에 O(가방 * 보석) 이라서 시간 초과가 발생했다.

 

문제 요구사항 보니까 1초에 N,K가 300000까지 가는거보면 O(NlogN) 안에서 끝을 봐야하는데... 

 

몇시간 머리를 싸맨 결과 아래와 같은 뇌피셜은 도출해냈다.

 

1) 가방과 보석을 무게 기준으로 오름차순 정렬

2) 처음 가방을 기준으로 해당 최대 무게에 해당하는 보석까지 하나씩 순서대로 정렬하면서 읽음 

3) 기준잡은 가방에 대해 들어갈 수 있는 맥시멈 보석의 무게까지 왔으면 거기서 가격 최댓값 뽑기

4) 다음 가방으로 가면서 전의 가방에서 정렬했던거 그대로 가져가면서 다음 가방의 최대 무게에 해당하는 보석까지 다시 하나씩 순서대로 정렬하면서 읽기

5) 3번과 마찬가지고 최댓값 뽑기 

 

이러면 선형 시간복잡도 안에서 해결이 될거 같았는데... 이걸 구현을 어떻게 할까에서 막혔었다.

 

우선순위 큐

 

전체가 아닌 어떤 선에서까지의 최대/최소를 구하고 싶으면 우선순위 큐에서 해당 선까지의 원소들만 넣고 큐에서 poll하면 그만이다. 

 

그리디에선 알고리즘 자체가 최대 / 최소를 뽑다보니 아무래도 우선순위 큐를 사용하는 경우가 많은 것 같다. 기억하자.