1202 : 보석 도둑 - Java

1202번: 보석 도둑

아이디어 생각해내는게 너무 어려웠다.
처음에는 보석 이 들어가길래 습관적 DP로 생각했는데, 각 가방의 크기가 너무 커서 배열로 담을 순 없다는 것을 깨달았다. 그래서 ‘가방의 사이즈를 배열의 값으로 담아 K개 저장한뒤 오름차순으로 정렬해야겠다’까지는 쉽게 생각이 났다. 그 뒤로 한 시간 이상 고민했고, 풀이방법만 봐야겠다 싶어 확인해봤는데 아이디어는 간단했다.

우선순위 큐를 두개 준비한다. 하나는 무게를 오름차순으로 정렬하고 하나는 가치를 기준으로 정렬할 것이다. 이를 각각 pqByWeight와 pqByValue라고 하겠다.

입력받은 모든 보석들을 pqByWeight에 넣는다.
가방을 입력받고 가방의 최대무게를 오름차순으로 정렬한다.

1. 각 가방의 최대 무게마다 해당 무게보다 낮은 원소들을 pqByWeight에서 꺼내 pqByValue에 넣어준다.
2. pqByValue에서 가장 첫 번째 원소 즉, 가장 가치가 높은 보석을 꺼내고 ans에 더해준다.
3. 마지막 가방까지 1 ~ 2를 반복한다.

이게 가능한 이유는 아까전 가방의 최대 무게를 오름차순으로 정렬했기 때문이다. 또한, pqByWeight에서 각 가방의 최대 무게까지만 원소를 pqByValue로 옮겨줬기 때문에 각 원소들은 늘 항상 현재 가방의 최대 무게보다 작음이 보장된다. 따라서, pqByValue에서 가장 가치가 큰 원소를 뽑아서 가방에 넣으면 무게는 현재가방보다 작고 가치가 가장 높은 원소들만 저장하게 된다.

다만, 보석이 가방에 넣을 수 있는 것보다 큰 경우나 모든 원소가 특정 하나의 무게보다 작은 경우에만 예외처리를 해주면 된다.

전체 코드

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static int N, K;

    static class Jewelry {
        int weight;
        int value;

        public Jewelry(int weight, int value) {
            this.weight = weight;
            this.value = value;
        }
    }

    static PriorityQueue<Jewelry> pqByWeight = new PriorityQueue<>((j1, j2) -> j1.weight - j2.weight);
    static PriorityQueue<Jewelry> pqByValue = new PriorityQueue<>((j1, j2) -> j2.value - j1.value);
    static int[] bagSize;
    static long ans;

    public static void main(String[] args) throws IOException {
//        System.setIn(Files.newInputStream(Paths.get("src/input.txt")));
//        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

        bagSize = new int[K];

        for (int i = 0; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int m = Integer.parseInt(st.nextToken()), v = Integer.parseInt(st.nextToken());
            pqByWeight.add(new Jewelry(m, v));
        }
        for (int i = 0; i < K; i++) {
            bagSize[i] = Integer.parseInt(br.readLine());
        }
        Arrays.sort(bagSize);

        int idx = 0;
        while (idx < K) {
            // 모든 원소가 특정 무게보다 작은 경우 제외
            while (!pqByWeight.isEmpty() && pqByWeight.peek().weight <= bagSize[idx]) {
                pqByValue.add(pqByWeight.poll());
            }
            Jewelry poll = pqByValue.poll();
            // 가방의 최대무게보다 보석 무게가 큰 경우
            if (poll != null) {
                ans += poll.value;
            }
            idx++;
        }
        bw.write(ans + "");
        bw.flush();
        bw.close();
        br.close();
    }


}