[백준 자바 1202] 보석 도둑 (골드2)

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

 

맨 아래에 전체 코드 있다. 

 

 

설명 

처음에 배낭문제 인줄 알았는데 가방이 여러 개라 어떻게 풀지 고민하다가 가방에 1개의 보석만 담을 수 있어서 아닌 것 같았다(?) 

 

일단 보석을 객체로 만들고, 

1. 가방 최대 무게 배열을 오름차순으로 정렬
2. 보석 리스트를 (1) 무게 오름차순, (2) 가격 내림차순 정렬 
3. 매 가방마다 보석 리스트 돌며 최대 무게 이하인 보석만 우선순위 큐에 넣기 (우선순위 큐는 2번과 같이 정렬됨)
4. 하나 poll 해와서 가격을 총 더하고, 더한 보석은 보석리스트에서 제거 출력 

 

인줄 알았는데 시간 초과가 났다. 아무래도 보석과 가방 개수가 최대 30만개에, 이중 반복문이면 900억이어서.. 

 

 

그래서 풀이 봤다! 

근데 풀이 설명이 좀 이해하기 힘들어서 최대한 내가 풀어서 설명을 해본다.. 

 

---

 

큰 흐름은 다음과 같다.

1. 가방 최대 무게 배열을 오름차순 정렬
2. 보석 우선순위큐를 (1)무게 오름차순 (2)가격 내림차순 기준으로 담기
3. 매 가방마다, 가방의 최대 무게 이하인 보석들을 보석 우선순위큐에서 모두 빼와서 가격 우선순위큐에 넣기
4. 가격 우선순위 큐에서 하나만 빼오고, 이를 가방마다 다 더해서 출력 

 

 

우선 보석 객체

static class Jewel implements Comparable<Jewel> {
    int weight, price;

    public Jewel(int weight, int price) {
        this.weight = weight;
        this.price = price;
    }

    public int compareTo(Jewel o) {
        if (this.weight != o.weight) {
            return Integer.compare(this.weight, o.weight);
        }

        return Integer.compare(o.price, this.price);
    }
}

 

무게 weight, 가격 price 를 두었고, Comparable<Jewel> 을 구현하였다. 

정렬 조건은 1. 무게 오름차순 2. 가격 내림차순. 

 

 

static int N, K;
static long answer;
static int[] bags;
static Queue<Jewel> jewels = new PriorityQueue<>();
static Queue<Integer> pricePQ = new PriorityQueue<>(Collections.reverseOrder());

 

그 다음 풀이에 쓰이는 변수들. 

보석 N개,

가방 K개,

총 보석 가격 합인 answer (보석 가격은 최대 100만, 보석 개수는 최대 30만이므로 long 타입이어야 한다), 

가방의 최대무게를 담은 배열 bags,

보석 우선순위큐 jewels,

보석의 가격만 담은 우선순위 큐 pricePQ (가격 내림차순이다) 

 

 

BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());

N = stoi(st.nextToken());
K = stoi(st.nextToken());

bags = new int[K];

for (int i = 0; i < N; i++) {
    st = new StringTokenizer(br.readLine());
    int weight = stoi(st.nextToken());
    int price = stoi(st.nextToken());

    jewels.add(new Jewel(weight, price));
}

for (int i = 0; i < K; i++) {
    bags[i] = stoi(br.readLine());
}

여기는 입력이다. stoi는 내가 커스텀으로 만든 함수인데, Integer.parseInt(string) 이다. 

 

Arrays.sort(bags);

for (int maxBagWeight : bags) {

    while (!jewels.isEmpty() && jewels.peek().weight <= maxBagWeight) {
        pricePQ.add(jewels.poll().price);
    }

    if (!pricePQ.isEmpty()) {
        answer += pricePQ.poll();
    }
}

System.out.println(answer);

핵심 풀이 부분. 

 

우선 가방 무게 배열을 오름차순 정렬을 해주고, 

 

보석 우선순위큐가 비어있지 않고, 피크 보석의 무게가 최대 가방 무게 이하라면 이를 빼와서 가격 우선순위큐에 가격을 넣어준다. 

 

그리고 가격 우선순위큐가 비어있지 않다면 하나 빼와서 answer에 더해준다. 

 

가격 우선순위큐에는 다음 가방을 처리할 때도 초기화 되지 않고 남아있다. 

 

핵심은 현재 무게가 허락하는 한 모든 보석들의 가격을 가져온 뒤, 매 가방마다 가장 높은 가격의 보석을 넣어준다는 것이다. 

 

 

 

---

전체 코드

 

 

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

public class Main {
    static int stoi(String s) {
        return Integer.parseInt(s);
    }

    static int N, K;
    static long answer;
    static int[] bags;
    static Queue<Jewel> jewels = new PriorityQueue<>();
    static Queue<Integer> pricePQ = new PriorityQueue<>(Collections.reverseOrder());

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        N = stoi(st.nextToken());
        K = stoi(st.nextToken());

        bags = new int[K];

        for (int i = 0; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int weight = stoi(st.nextToken());
            int price = stoi(st.nextToken());

            jewels.add(new Jewel(weight, price));
        }

        for (int i = 0; i < K; i++) {
            bags[i] = stoi(br.readLine());
        }

        Arrays.sort(bags);

        for (int maxBagWeight : bags) {

            while (!jewels.isEmpty() && jewels.peek().weight <= maxBagWeight) {
                pricePQ.add(jewels.poll().price);
            }

            if (!pricePQ.isEmpty()) {
                answer += pricePQ.poll();
            }
        }

        System.out.println(answer);
    }

    static class Jewel implements Comparable<Jewel> {
        int weight, price;

        public Jewel(int weight, int price) {
            this.weight = weight;
            this.price = price;
        }

        public int compareTo(Jewel o) {
            if (this.weight != o.weight) {
                return Integer.compare(this.weight, o.weight);
            }

            return Integer.compare(o.price, this.price);
        }
    }
}