문제 출처

단계별로 풀어보기 > 우선순위 큐 > 절댓값 힙
https://www.acmicpc.net/problem/11286

문제 설명

N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다.
만약 x가 0이라면 배열에서 절댓값이 가장 작은 값을 출력하고 그 값을 배열에서 제거하라.
x가 0이 아니라면 배열에 x를 추가하라.
단, 배열이 비어있는 경우 x가 0이라면 0을 출력하라.

접근 방법

PriorityQueue를 생성할 때, 조건을 설정한다.
두 개의 절댓값이 같으면 return o1 - o2를 하도록, 절댓값이 다르면 절댓값이 작은 수가 먼저오게 return Math.abs(o1) - Math.abs(o2)로 조건을 설정한다.
이 후, x가 0이면 PriorityQueue에서 절댓값이 가장 작은 수를 빼내어 출력하도록, 0인 경우에 PriorityQueue가 비어있으면 0을 출력하도록, 마지막으로 x가 0이 아닌 경우 PriorityQueue에 x를 삽입하도록 한다.

import java.io.*;
import java.util.Collections;
import java.util.PriorityQueue;

public class 절대값_힙 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        int N = Integer.parseInt(br.readLine());

        PriorityQueue<Integer> pq = new PriorityQueue<>((o1,o2) ->{
            if(Math.abs(o1) == Math.abs(o2)){
                return o1 - o2;
            }
            return Math.abs(o1) - Math.abs(o2);
        });

        StringBuilder sb = new StringBuilder();
        for(int i = 0; i < N; i++){
            int x = Integer.parseInt(br.readLine());

            if(x == 0){
                if(pq.size() == 0){
                    sb.append(0).append("\n");
                    continue;
                }

                sb.append(pq.poll()).append("\n");
            }else{
                pq.offer(x);
            }
        }

        bw.write(sb.toString());
        bw.flush();
        bw.close();
        br.close();
    }
}

알게된 점

처음 풀이할 때는, PriorityQueue에 조건을 넣지 않고, 단순히 2개를 이용해야겠다는 생각으로 접근했으나, 문제에서 주어진 예시 입력값에 대한 결과도 틀리고, 시간 복잡도가 remove 때문에 O(N^2)까지 늘어나서, 효율성 부분에서도 좋지 않았다.

import java.io.*;
import java.util.Collections;
import java.util.PriorityQueue;

public class 절대값_힙 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        PriorityQueue<Integer> minQ = new PriorityQueue<>(Collections.reverseOrder());
        PriorityQueue<Integer> maxQ = new PriorityQueue<>();

        int N = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();

        for (int i = 0; i < N; i++) {
            int x = Integer.parseInt(br.readLine());

            if (x == 0) {
                if (minQ.size() == 0) {
                    sb.append(0).append("\n");
                    continue;
                }

                int min = minQ.peek();
                int max = maxQ.peek();

                if (Math.abs(max) < Math.abs(min)) {
                    int k = maxQ.poll();
                    sb.append(k).append("\n");
                    minQ.remove(k);
                } else {
                    int k = minQ.poll();
                    sb.append(k).append("\n");
                    maxQ.remove(k);
                }
            } else {
                minQ.offer(x);
                maxQ.offer(x);
            }
        }

        bw.write(sb.toString());
        bw.flush();
        bw.close();
        br.close();
    }
}

정답에 대한 코드의 시간복잡도는 O(NlogN)이다.

문제푼 흔적