[실압코] 우선순위 큐 (priority_queue)

들어가며

 오늘은 우선순위 큐에 대해서 알아보겠습니다. 먼저 우선순위 큐의 개념에 대해서 소개하고, 이를 활용하는 다익스트라 알고리즘을 활용하는 문제까지 한번 풀어보겠습니다.

우선순위 큐

 우선순위 큐란 힙 형태의 자료구조로 우선순위가 높은 원소를 먼저 내보내는 자료구조입니다. 그림으로 보겠습니다. 다음은 크기가 클수록 우선순위가 높은 우선순위 큐입니다.

출처: https://www.geeksforgeeks.org/dsa/priority-queue-set-1-introduction/

 완전 이진트리의 형태로 구현되어 있기 때문에 가장 우선순위가 높은 원소를 뽑는 top()은 O(1)의 시간, 원소를 넣고 빼는 push(), pop()은 O(logn)의 시간이 걸립니다. 삽입과 삭제가 빈번하고, 최우선값을 자주 사용해야 하는 자료구조 필요할 때 유용한 자료구조입니다.

priority_queue

 이번엔 C++에서 어떤 식으로 우선순위 큐를 구현하고 사용하는지 알아보겠습니다.

#include <queue>

using namespace std;

int main()
{
    // 기본적으로 최대값을 뽑을 수 있는 maxheap으로 설정
    priority_queue<int> pq;

    pq.push(10);    // 원소를 삽입
    pq.pop();       // 최상위 원소를 제거
    pq.top();       // 최상위 원소를 반환

    pq.size();      // 원소의 개수
    pq.empty();     // 비어있는지 여부
    
    return 0;
}

 

 우선순위 큐를 사용할 때 최소 힙으로 사용하고 싶다면 이런 방법들이 있습니다.

#include <vector>
#include <queue>

// greater<>
#include <functional>

using namespace std;

int main()
{
    // maxheap을 minheap처럼 사용하기
    priority_queue<int> pq;
    int value = 10;

    // 부호를 뒤집어 주면 최소값 -> 최대값이 되기 때문에
    // minheap처럼 사용할 수 있다.
    pq.push(-value);

    // 진짜 minheap 사용하기: 타입(int), 내부 컨테이너(vector<int>), 비교 함수(greater<int>) 순으로 작성
    priority_queue<int, vector<int>, greater<int>> minheap;

    return 0;
}

 

 이번에는 사용자 지정 비교함수를 구현해 보겠습니다. pair <int, int>는 원래 first가 클수록 우선순위가 높지만, second를 기준으로 우선순위를 정하도록 구현해 보겠습니다. 또한 우선순위 큐의 Compare를 작성할 때 주의할 점은 일반적인 배열을 sort 할 때 사용하는 람다 함수의 기준과는 반대로 작성해야 합니다. 일반적인 배열의 정렬 우선순위를 정의할 때 작은 값을 앞에 오게 하려면 a < b의 형태로 작성하지만 우선순위 큐에서 작은값을 앞에 오게 하고 싶다면 a > b 의 형태로 작성해야 합니다.

#include <vector>
#include <queue>

using namespace std;

struct Compare
{
    bool operator()(const pair<int, int>& a, const pair<int, int>& b)
    {
        // pair의 second가 클수록 우선순위가 높은 형태
        return a.second < b.second;
    }
};

int main()
{
    // 사용자 지정 우선순위
    priority_queue<pair<int, int>, vector<pair<int, int>>, Compare> pq;

    pq.push({ 10, 20 });
    pq.push({ 20, 10 });

    // 10 20
    cout << pq.top().first << ' ' << pq.top().second;

    return 0;
}

 

BOJ 1854. K번째 최단경로 찾기

 다익스트라를 활용하여 K번째 최단 경로를 찾는 문제입니다. 다익스트라 알고리즘은 여기에서 설명했습니다. 일반적인 다익스트라는 특정지점 a에서 b까지 가는 최단경로를 구하지만 최단경로를 구한 후에도 계속해서 진행한다면 K번째 최단경로를 찾을 수 있다는 것이 핵심 아이디어입니다.

// BOJ 1854. K번째 최단경로 찾기
#include <iostream>
#include <vector>
#include <queue>

using namespace std;

int main()
{
    // 빠른 입출력
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m, k;
    cin >> n >> m >> k;

    // 인접 리스트
    vector<vector<pair<int, int>>> adj(n + 1);
    for (int i = 0; i < m; i++)
    {
        int a, b, c;
        cin >> a >> b >> c;

        adj[a].push_back({ b, c });
    }

    // 몇번째 방문인지
    vector<int> visited(n + 1);

    // 정답 배열
    vector<int> answer(n + 1, -1);

    // 우선순위 큐
    priority_queue<pair<int, int>> pq;
    pq.push({ 0, 1 });

    // 다익스트라
    while (!pq.empty())
    {
        auto [cost, node] = pq.top();
        pq.pop();

        // 이미 k번 이상 방문했다면 건너뜀
        if (visited[node] >= k)
        {
            continue;
        }

        visited[node]++;

        // 이번이 k번째 방문이라면 정답 배열에 기록
        if (visited[node] == k)
        {
            answer[node] = -cost;
        }

        for (auto [b, c] : adj[node])
        {
            // 이미 k번 이상 방문했다면 건너뜀
            if (visited[b] >= k)
            {
                continue;
            }

            pq.push({ cost - c, b });
        }
    }

    // 정답 출력
    for (int i = 1; i <= n; i++)
    {
        cout << answer[i] << '\n';
    }

    return 0;
}

 

 다익스트라 알고리즘의 특성상 K번째로 특정 노드에 방문한다면 해당 경로가 K번째 최단경로가 되기 때문에 각 노드에 방문한 횟수를 기록하고 K번째가 되는 순간 정답을 기록하면 됩니다. 또한 특정 노드에 K+1번째 이상 도착했다면 해당 경로를 포함한 모든 경로는 K번째 최단경로가 될 수 없습니다. 때문에 이를 가지치기해주어 최적화할 수 있습니다.

마치며

 오늘은 우선순위 큐와 다익스트라를 활용하는 문제를 풀어보았습니다. 우선순위 큐의 개념은 간단하고 유용하기 때문에 알고리즘 문제를 풀 때 많은 도움이 됩니다. 감사합니다.