[백준 - 23022번] 숙제 - Java

문제 링크

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

문제

티어: 골드 2
알고리즘: 그리디, 우선순위 큐, 정렬

입력

첫 줄에 테스트 케이스의 수 T가 주어진다.

각 테스트 케이스는 세 줄에 걸쳐 주어진다.

첫 줄에 두 정수 n과 S가 공백으로 구분되어 주어진다.

둘째 줄에 숙제가 언제 나오는지 나타내는 n개의 정수가 (t[1], ..., t[n]) 공백으로 구분되어 주어진다.

셋째 줄에 숙제의 벌점을 나타내는 n개의 정수가 (v[1], ..., v[n]) 공백으로 구분되어 주어진다.

출력

각 테스트 케이스의 정답인 최소 벌점을 각 줄에 출력한다.

제한

• 1 ≤ T ≤ 10
• 1 ≤ n ≤ 100,000
• 1 ≤ S ≤ 2,000,000,000
• 1 ≤ t[i] ≤ 2,000,000,000
• 1 ≤ v[i] ≤ 40,000

풀이

Albert가 모든 숙제를 다 제출하면서 달성 가능한 최소한의 벌점을 구해야 한다.

직감적으로 최소한의 벌점을 위해서는 v 값이 큰 순으로 가능한 숙제를 제출해야 할 것 같다는 느낌이 든다.

이 직감은 맞지만, 당연히 증명이 되어야 한다.

입력이 다음과 같을 때
5 3
1 2 3 4 5
8 3 2 13 3

S=3일 때 풀 수 있는 숙제는 [1,8], [2,3], [3,2]이 된다.
현재 시간에서 v가 가장 큰 1번 숙제를 푸는게 최선의 선택일까?
그렇다면 t는 선택하는데 영향을 끼치지 않을까?

결론부터 말하자면 1번을 푸는 것은 최선의 선택이며, t는 영향을 끼치지 않는다. 생각해 보면 1초 지났을 때 증가하는 값은 v다.

그러니까 S=4일 때 1번 숙제는 8 증가하고, 2번 숙제는 3 증가하고, 3번 숙제는 2 증가한다.

이 증가하는 값을 최소로 하려면 t와는 상관없이 현재 풀 수 있는 숙제 중 v가 가장 큰 숙제를 제출하면 된다.

이를 구현하기 위해서 우선순위 큐가 필요하다.

처음 S를 기준으로 풀 수 있는 숙제를 우선순위 큐(V를 기준으로 내림차순)에 넣고, poll해 준다. 제출할 때마다 S를 1씩 증가시키며 또 풀 수 있는 숙제를 우선순위 큐에 넣고 poll해준다.

여기서 주의할 점은 S는 최대 20억이기 때문에 계속 1씩 증가시키면 당연히 TLE가 난다. 그래서 우선순위 큐에 아무런 값도 없을 때는 굳이 S를 1씩 증가시킬 필요가 없기 때문에 S를 풀지 않은 숙제중 가장 빨리 풀 수 있는 숙제의 t로 변환해주고, 가능한 숙제를 우선순위 큐에 넣어줘야 한다. (가능한 숙제를 찾을 때는 숙제들이 시간 순으로 정렬된 배열이 필요함)

이 풀이의 시간 복잡도는 O(N * logN)이다.

소스 코드

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

class Homework implements Comparable<Homework> {
    int t, v;
    Homework(int t, int v) {
        this.t = t;
        this.v = v;
    }
    
    @Override
    public int compareTo(Homework o) {
        if(this.v < o.v) {
            return 1;
        } else if(this.v > o.v) {
            return -1;
        }
        return 0;
    }
}

public class Main {
    static int T;
  public static void main(String args[]) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      T = Integer.parseInt(br.readLine());
      StringBuilder sb = new StringBuilder();
      for(int k=0; k<T; k++) {
          StringTokenizer st = new StringTokenizer(br.readLine());
          int n = Integer.parseInt(st.nextToken());
          int S = Integer.parseInt(st.nextToken());
          Homework[] hm = new Homework[n];
          StringTokenizer st2 = new StringTokenizer(br.readLine());
          StringTokenizer st3 = new StringTokenizer(br.readLine());
          for(int i=0; i<n; i++) {
              hm[i] = new Homework(Integer.parseInt(st2.nextToken()), Integer.parseInt(st3.nextToken()));
          }
          
          Arrays.sort(hm, new Comparator<Homework>() {
             @Override
             public int compare(Homework h1, Homework h2) {
                 return Integer.compare(h1.t, h2.t);
             }
          });
          
          PriorityQueue<Homework> pque = new PriorityQueue<>();
          long answer = 0;
          int cursor = 0;
          for(int i=0; i<n; i++) {
              if(hm[i].t > S) {
                  break;
              }
              cursor += 1;
              pque.add(hm[i]);
          }
          while(true) {
              if(pque.isEmpty()) {
                  if(cursor == n) {
                      break;
                  }
                  S = hm[cursor].t;
                  for(int i=cursor; i<n; i++) {
                      if(hm[i].t > S) {
                          break;
                      }
                      cursor += 1;
                      pque.add(hm[i]);
                  }
              }
              Homework h = pque.poll();
              answer += (long) (S - h.t) * (long) h.v;
              S += 1;
              
              if(!pque.isEmpty()) {
                  for(int i=cursor; i<n; i++) {
                      if(hm[i].t > S) {
                          break;
                      }
                      cursor += 1;
                      pque.add(hm[i]);
                  }
              }
              
          }
          
          sb.append(answer).append("\n");
      }
      System.out.println(sb.toString().trim());
  }
}