[백준 - 22899번] 오렌지컵 출제하기 - Java

문제 링크

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

문제

티어: 골드 2
알고리즘: 그리디, 우선순위 큐, 해시맵

입력

첫 번째 줄에는 문제 아이디어의 수 $N$과 출제할 문제의 수 $K$가 주어진다.

두 번째 줄에는 각 문제의 출제자의 번호 $a_1, a_2, \cdots, a_N$이 주어진다.

세 번째 줄에는 각 문제의 준비 시간 $b_1, b_2, \cdots, b_N$이 주어진다.

출력

$L$이 $1$부터 $N$까지의 정수일 때 문제의 답을 띄어쓰기로 구분하여 출력한다. 단, 조건을 만족하도록 문제 출제가 불가능한 경우 -1을 출력한다.

풀이

bi의 합의 최솟값을 구해야 한다.

최솟값은 당연히 bi가 작은 것을 우선적으로 선택하면 되는데, 출제자마다 L개의 문제만 낼 수 있기 때문에 그렇게 간단하진 않다.

내가 풀이한 방식은 다음과 같다.

각 출제자마다 우선순위 큐를 두고 문제의 bi를 넣는다. (오름차순 정렬)

그리고 L을 1부터 N까지 반복하면서, 각 출제자마다 우선순위 큐에서 값을 꺼내고, 내림차순으로 정렬하는 우선순위 큐(problems)에 넣는다. 그리고 그 값은 sum에 계속 더해준다.

그러면 조건(L개의 문제만 낼 수 있는)에는 부합되는데, 여기서 size가 K개보다 작다면 -1를 출력하고, K개보다 크다면 size - K개 만큼 poll() 해주면 된다. 그리고 poll()한 값은 당연히 sum에서 빼주면, 현재 L에서 최적 해를 구할 수 있다.

전체 문제가 어차피 N개이기 때문에 L을 1부터 N까지 반복한다고 해도 시간 복잡도는 O(N * logN)이 된다.

그런데 출제자마다 우선순위 큐를 정의해 줄 때 적절한 자료구조를 사용해야 한다. 만약 무작정 1~N까지의 출제자의 우선순위 큐를 정의한다면 O(N^2)으로 시간 초과가 발생한다.

그래서 해시맵을 사용해야 한다. 해시맵을 사용한다면, 현재 문제를 더 낼 수 있는 출제자만큼 반복을 돌기 때문에 O(N * logN)을 유지할 수 있다.

소스 코드

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

public class Main {
    static int N, K;
  public static void main(String args[]) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      StringTokenizer st = new StringTokenizer(br.readLine());
      N = Integer.parseInt(st.nextToken());
      K = Integer.parseInt(st.nextToken());
      
      HashMap<Integer, PriorityQueue<Integer>> tester = new HashMap<>();
      StringTokenizer st1 = new StringTokenizer(br.readLine());
      StringTokenizer st2 = new StringTokenizer(br.readLine());
      
      for(int i=0; i<N; i++) {
          int a = Integer.parseInt(st1.nextToken());
          int b = Integer.parseInt(st2.nextToken());
          if(tester.get(a) == null) {
              tester.put(a, new PriorityQueue<>());
          }
          tester.get(a).add(b);
      }
      
      PriorityQueue<Long> problems = new PriorityQueue<>(Comparator.reverseOrder());
      StringBuilder sb = new StringBuilder();
      long sum = 0;
      for(int i=1; i<=N; i++) {
          Iterator<Map.Entry<Integer, PriorityQueue<Integer>>> iterator = tester.entrySet().iterator();
          while(iterator.hasNext()) {
              Map.Entry<Integer, PriorityQueue<Integer>> entry = iterator.next();
              long b = (long) entry.getValue().poll();
              problems.add(b);
              sum += b;
              if(entry.getValue().isEmpty()) {
                  iterator.remove();
              }
          }
          if(problems.size() >= K) {
              int size = problems.size();
              for(int j=1; j<=size - K; j++) {
                  sum -= problems.poll();
              }
              sb.append(sum).append(" ");
          } else {
              sb.append(-1).append(" ");
          }
      }
      System.out.println(sb.toString().trim());
  }
}