문제

Programmers, 이중우선순위큐

핵심

• 입력의 크기가 $10^6$이라 대략 $O(N\times\log_2N)$ 이하의 알고리즘을 사용한다.
• 큐에서 주어진 숫자를 삽입하고, 최댓값과 최솟값을 삭제하는 연산이 주어진다. 삭제하는 연산에서 효율적으로 최솟값이나 최댓값을 찾아내는 게 핵심이다. 모든 연산을 처리한 뒤에 큐가 비어있지 않다면 최댓값, 최솟값을 반환한다.
• C++의 multiset을 사용하거나 Java의 TreeMap을 조금 변형하여 중복을 담을 수 있는 이진 검색 트리 자료구조를 만들면 직관적으로 해결이 가능하다. 최댓값은 정렬된 이진 검색 트리의 끝에서 지우고, 최솟값은 제일 앞에서 지우면 되기 때문이다.
• 힙 자료구조에 익숙해지기 위해 우선순위 큐를 이용한다. 우선순위 큐는 최대 힙 또는 최소 힙으로 구현할 수 있는데, 최소 힙으로 구현할 경우 최솟값은 알 수 있지만 최댓값은 알 수 없다. 역으로도 마찬가지다. 그렇다면 최대 힙 우선순위 큐와 최소 힙 우선순위 큐 2개가 필요하다.
• 최댓값을 저장하기 위한 우선순위 큐에선 최댓값을 제거하는 연산만 처리한다. 그 이유는 최대 힙 우선순위 큐에서는 최솟값을 한 번에 찾아서 제거하는 기능이 없기 때문이다. 만약 큐에 10, 5를 넣고 최솟값을 제거하는 연산을 3번 했다면 최대 힙 우선순위 큐에서는 아무런 연산을 하지 않는다. 뭔가 이상하지 않은가? 최대 힙에서도 최솟값이 지워졌다는 걸 반영해 주어야 한다. 이를 위해 큐의 크기를 관리하는 변수를 선언한다. 최솟값을 제거할 때 최대 힙 우선순위 큐에서는 아무런 작업도 하지 않지만, 만약 모든 원소를 지웠다면 이를 반영하기 위해 큐의 크기가 0이 되면 모든 큐를 비워주는 작업을 한다. 이는 최소 힙 우선순위 큐에서도 마찬가지다.

개선

• C++ 우선순위 큐 자료구조에서는 clear 연산을 지원하지 않는다. 하지만 swap() 메서드를 이용하면 매우 효율적으로 구현할 수 있다. 핵심은 비어 있는 큐의 메모리 주소와 가득 담겨 있는 큐의 메모리 주소만 바꿔치기 하는 것이고, STL 특성상 참조 범위를 넘어서면 자동으로 메모리 해제가 되기 때문에 깔끔한 방식이다. 하지만 우선순위 큐가 커스텀 비교 함수를 사용하고 있을 때 추가해야 할 코드 작업이 늘어나기 때문에 시간이 급박한 코딩테스트에서는 swap 연산까지 처리할 필요는 없어 보인다.

코드

시간복잡도

• $O(N)$

C++

#include <string>
#include <vector>
#include <queue>
#include <iostream>
#include <string>
using namespace std;

vector<int> solution(vector<string> operations) {
    priority_queue<int> maxHeap;
    auto comp = [](auto& a, auto& b) { return a > b; };
    priority_queue<int, vector<int>, decltype(comp)> minHeap(comp);
    int pqSize = 0;
    for (auto str : operations) {
        char cmd = str[0];
        int num = stoi(str.erase(0, 2));
        if (cmd == 'I') {
            maxHeap.push(num);
            minHeap.push(num);
            ++pqSize;
        }
        else {
            if (pqSize > 0) {
                if (num == 1)
                    maxHeap.pop();
                else
                    minHeap.pop();         
                --pqSize;
            }
        } 
        if (pqSize == 0) {
            while (!maxHeap.empty())
                maxHeap.pop();
            while (!minHeap.empty())
                minHeap.pop();
        }
    }
    vector<int> answer = {0, 0};
    if (pqSize > 0)
        answer = {maxHeap.top(), minHeap.top()};
    return answer;
}

Java

import java.util.*;

class Solution {
    public int[] solution(String[] operations) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b.compareTo(a));
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        int pqSize = 0;
        for (var str : operations) {
            char cmd = str.charAt(0);
            int num = Integer.parseInt(str.substring(2));
            if (cmd == 'I') {
                maxHeap.add(num);
                minHeap.add(num);
                ++pqSize;
            } else {
                if (pqSize > 0) {
                    if (num == 1)
                        maxHeap.poll();
                    else
                        minHeap.poll();
                    --pqSize;
                }
            }
            if (pqSize == 0) {
                maxHeap.clear();
                minHeap.clear();
            }
        }
        int[] answer = {0, 0};
        if (pqSize > 0) {
            answer[0] = maxHeap.peek();
            answer[1] = minHeap.peek();
        }
        return answer;
    }
}