[Java] 백준 BOJ / 7662번: 이중 우선순위 큐

이중 우선순위 큐

문제

---

문제 설명

테스트케이스 T 와 각 테스트케이스에 대해서 들어올 입력의 개수 k 가 주어진다.

입력은 어떤 연산을 할지, D 또는 I 가 들어오며, I 는 삽입 연산을 의미한다.

D 는 삭제를 뜻하는데 다음으로 들어오는 숫자가 1 이라면 최댓값을, -1 이라면 최솟값을 삭제하면 된다.

유의할 점은 만약에 최댓값이나 최솟값이 여러개 존재하는 경우 하나만 삭제하면 된다.

---

구현

처음 구현한 방식은 문제의 제목에 맞게 우선순위 큐를 두개를 두었다.

import static java.util.Collections.reverse;

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

public class Main_7662 {
  static int T;
  static int k;
  public static void main(String[] args) throws IOException {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    T = Integer.parseInt(br.readLine());

    while (T-- > 0) {
      StringBuilder sb = new StringBuilder();

      k = Integer.parseInt(br.readLine());
      //오름차순 정렬
      PriorityQueue<Integer> maxHeap = new PriorityQueue<>();
      //내림차순 정렬
      PriorityQueue<Integer> minHeap = new PriorityQueue<>(Collections.reverseOrder());

      for(int i = 0; i < k; i++){
        StringTokenizer st = new StringTokenizer(br.readLine());
        String cmd = st.nextToken();
        int num = Integer.parseInt(st.nextToken());
        switch(cmd){
          case "I" :
            maxHeap.add(num);
            minHeap.clear();
            minHeap.addAll(maxHeap);
            break;
          case "D" :
            //최댓값 삭제하고 내림차순으로 바꿔서 minHeap 초기화
            if (num == 1) {
              minHeap.poll();
              maxHeap.clear();
              maxHeap.addAll(minHeap);
            }
            //최솟값 삭제
            else {
              maxHeap.poll();
              minHeap.clear();
              minHeap.addAll(maxHeap);
            }
            break;
        }
      }

      if(maxHeap.isEmpty()){
        sb.append("EMPTY");
      }
      else{
        sb.append(minHeap.poll()).append(" ").append(maxHeap.poll());
      }
      System.out.println(sb);
    }
  }
}

이 코드와 같이 오름차순으로 정렬할 maxHeap 과 내림차순으로 정렬된 상태를 유지하는 minHeap 을 두고 연산이 들어올때마다 전부 초기화해주는 방법을 사용했다.

이 경우 k 가 10^6 이라면 10^6 * log(10^6) 으로 많은 시간을 소요하기 때문에 다른 방법을 찾다가 TreeMap 이라는 자료구조를 사용하여 문제를 해결할 수 있다는 것을 알았다.

TreeMap 이란 정렬된 상태로 저장되며 최댓값과 최솟값에 접근을 할 수 있는 장점이 있다.

따라서 I 연산이 들어온다면 아래와 같이 숫자를 넣어주고, 개수를 늘려준다.

case "I" :
	map.put(num, map.getOrDefault(num, 0) + 1);
	break;

이 함수는 num 이라는 숫자가 존재하면, 그 값에 1을 더하고, 존재하지 않는다면 기본값인 0에서 1을 더해서 저장하는 의미이다.

그리고 D 연산에서는 map 이 비어있는지 확인하고, 비어있지 않다면 다음 연산자를 확인하고, 최댓값 또는 최솟값을 지워주면 된다.

case "D" :
	if (map.isEmpty()) {
		continue;
	}
	long key = (num == 1) ? map.lastKey() : map.firstKey();
	int count = map.get(key);

	if (count == 1) {
		map.remove(key);
	} else{
		map.put(key, count - 1);
	}
    break;

---

코드

import static java.util.Collections.reverse;

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

public class Main_7662 {
  static int T;
  static int k;
  public static void main(String[] args) throws IOException {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    T = Integer.parseInt(br.readLine());

    while (T-- > 0) {
      StringBuilder sb = new StringBuilder();

      k = Integer.parseInt(br.readLine());
      //오름차순 정렬
      TreeMap<Long, Integer> map = new TreeMap<>();

      for(int i = 0; i < k; i++){
        StringTokenizer st = new StringTokenizer(br.readLine());
        String cmd = st.nextToken();
        long num = Long.parseLong(st.nextToken());
        switch(cmd){
          case "I" :
            map.put(num, map.getOrDefault(num, 0) + 1);
            break;
          case "D" :
            if (map.isEmpty()) {
              continue;
            }
            long key = (num == 1) ? map.lastKey() : map.firstKey();
            int count = map.get(key);

            if (count == 1) {
              map.remove(key);
            } else{
              map.put(key, count - 1);
            }
            break;
        }
      }
      if(map.isEmpty()){
        sb.append("EMPTY");
      }
      else{
        sb.append(map.lastKey()).append(" ").append(map.firstKey());
      }
      System.out.println(sb);
    }
  }
}