[Programmers] #42628 - 이중우선순위큐 [Java][C++][Python]

문제 링크

---

1. 아이디어

이중 우선순위 큐라는 자료구조를 만드는 문제로 큐에 주어진 숫자를 삽입하는 연산, 큐에서 최댓값을 제거하는 연산, 큐에서 최솟값을 제거하는 연산을 빠르게 수행할 수 있어야 한다.

기본 아이디어는 최소힙, 최대힙 두 개와 각 원소의 제거 여부를 표시하는 방문 체크 배열을 활용하는 것이다. 큐에 주어진 숫자를 삽입하는 연산에서는 두 힙 모두에 삽입을 하고, 큐에서 최댓값을 제거하는 연산은 최대힙을, 큐에서 최솟값을 제거하는 연산은 최소힙을 선택해 원소를 제거한 후 제거 여부를 마킹하면 된다. 이러면 반대쪽 큐에는 아직 제거해야 했을 숫자가 남아있지만 제거 여부를 마킹으로 알 수 있으므로 큐에서 제거하기 전에 마킹된 값이면 이미 제거했어야 했을 값이므로 빼고 다시 제거할 값을 찾는 방식으로 이중 우선순위 큐를 구현할 수 있다.

---

2. 복잡도

접근	시간	공간
풀이	$O(N \log N)$	$O(N)$

($N$ = operations의 길이)

---

3. 코드

풀이 [Java][C++][Python]

import java.util.*;

class Solution {
    public int[] solution(String[] operations) {
        TreeMap<Integer, Integer> map = new TreeMap<>();

        for (String op : operations) {
            String[] parts = op.split(" ");
            int x = Integer.parseInt(parts[1]);

            if (parts[0].equals("I")) {
                map.put(x, map.getOrDefault(x, 0) + 1);
            } else {
                if (map.isEmpty()) continue;

                int key = (x == 1) ? map.lastKey() : map.firstKey();
                int cnt = map.get(key);
                if (cnt == 1) {
                    map.remove(key);
                } else {
                    map.put(key, cnt - 1);
                }
            }
        }

        if (map.isEmpty()) return new int[]{0, 0};
        return new int[]{map.lastKey(), map.firstKey()};
    }
}

Java는 TreeMap을 쓰면 이중 우선순위 큐를 직접 구현하지 않아도 된다. 키가 정렬된 채로 유지되어 최솟값과 최댓값을 firstKey, lastKey로 바로 조회할 수 있다. 같은 값이 여러 번 들어올 수 있어 값을 키, 등장 횟수를 값으로 담았고, 삭제할 때 횟수가 1이면 키를 지우고 아니면 횟수만 줄였다.

#include <bits/stdc++.h>
using namespace std;

vector<int> solution(vector<string> operations) {
    multiset<int> ms;

    for (string& op : operations) {
        int x = stoi(op.substr(2));

        if (op[0] == 'I') {
            ms.insert(x);
        } else {
            if (ms.empty()) continue;
            ms.erase(x == 1 ? prev(ms.end()) : ms.begin());
        }
    }

    if (ms.empty()) return {0, 0};
    return {*ms.rbegin(), *ms.begin()};
}

C++는 std::multiset이 중복을 허용하며 정렬된 순서로 저장하므로 이중 우선순위 큐의 로직을 그대로 적용할 수 있는 자료구조다.

import heapq


def solution(operations):
    min_h, max_h = [], []
    erased = [False] * len(operations)

    for i, op in enumerate(operations):
        cmd, x = op.split()
        x = int(x)

        if cmd == "I":
            heapq.heappush(min_h, (x, i))
            heapq.heappush(max_h, (-x, i))
        else:
            h = max_h if x == 1 else min_h
            while h and erased[h[0][1]]:
                heapq.heappop(h)

            if not h:
                continue

            erased[h[0][1]] = True
            heapq.heappop(h)

    while min_h and erased[min_h[0][1]]:
        heapq.heappop(min_h)
    while max_h and erased[max_h[0][1]]:
        heapq.heappop(max_h)

    if not min_h or not max_h:
        return [0, 0]
    return [-max_h[0][0], min_h[0][0]]

---