C++ 프로그램: q개의 쿼리에 대해 그래프에서 정점 k를 경유하는 최단 비용 경로 찾기

n개의 정점으로 구성되어 있고 최소한으로만 연결된(즉, 트리 형태의) 그래프가 주어졌다고 가정해 보겠습니다. 간선은 {출발점, 도착점, 가중치} 형식의 배열로 제공되며, 여기에 {출발점, 도착점} 형식의 쿼리가 q개 주어집니다. 각 쿼리마다 출발점에서 정점 k를 경유하여 도착점에 도달하는 최단 비용 경로를 찾고, 해당 경로의 비용을 출력해야 합니다.

예를 들어 입력이 n = 6, q = 3, k = 1, edges = {{1, 2, 2}, {1, 3, 4}, {3, 4, 2}, {3, 5, 3}, {5, 6, 2}}, queries = {{1, 4}, {2, 6}, {2, 5}}라면 출력은 6, 11, 9가 됩니다.

핵심 아이디어

이 문제에서 중요한 조건은 그래프가 '최소 연결(minimally connected)', 즉 트리라는 점입니다. 트리에서는 임의의 두 정점 사이의 경로가 항상 유일하므로, 출발점과 도착점이 정점 k를 기준으로 서로 다른 서브트리에 속해 있다면 그 경로는 반드시 k를 지나게 됩니다. 따라서 정점 k에서 시작하는 DFS를 한 번만 수행하여 모든 정점까지의 거리를 미리 계산해 두면, 각 쿼리의 답은 dist(x, k) + dist(y, k)로 O(1) 시간 안에 구할 수 있습니다.

풀이 단계

이 문제는 다음 단계를 따라 해결할 수 있습니다.

크기가 n * n인 pair의 2차원 배열 graph를 정의합니다
크기가 n인 배열 pathTotal을 정의합니다
dfs() 함수를 정의합니다. 이 함수는 a, b를 매개변수로 받습니다
    graph[a]의 각 값 i에 대해:
        i의 첫 번째 값이 b와 같다면:
            아래 부분을 건너뛰고 다음 반복으로 넘어갑니다
        pathTotal[i의 첫 번째 값] := pathTotal[a] + i의 두 번째 값
        dfs(i의 첫 번째 값, a)
i := 0으로 초기화하고, i < n - 1인 동안 i를 1씩 증가시키며 반복합니다:
    a := edges[i]의 첫 번째 값
    b := edges[i]의 두 번째 값
    c := edges[i]의 세 번째 값
    a와 b를 각각 1 감소시킵니다
    graph[a] 끝에 pair (b, c)를 삽입합니다
    graph[b] 끝에 pair (a, c)를 삽입합니다
k를 1 감소시킵니다
dfs(k, k)를 호출합니다
i := 0으로 초기화하고, i < q인 동안 i를 1씩 증가시키며 반복합니다:
    x := queries[i]의 첫 번째 값
    y := queries[i]의 두 번째 값
    x와 y를 각각 1 감소시킵니다
    pathTotal[x] + pathTotal[y]를 출력합니다

구현 예제

다음 C++ 구현 예제를 통해 더 자세히 이해해 보겠습니다.

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

vector<vector<pair<int,int>>> graph;
vector<int> pathTotal;
int k;

void dfs(int a, int b){
    for(auto i : graph.at(a)){
        if(i.first == b) continue;
        pathTotal.at(i.first) = pathTotal.at(a) + i.second;
        dfs(i.first,a);
    }
}
void solve(int n, int q, vector<tuple<int, int, int>> edges,
vector<pair<int, int>> queries){
    int a, b, c, x, y;
    graph.resize(n);
    pathTotal.resize(n);
    for(int i = 0; i < n - 1; i++){
        a = get<0> (edges[i]);
        b = get<1> (edges[i]);
        c = get<2> (edges[i]);
        a--, b--;
        graph.at(a).push_back(make_pair(b, c));
        graph.at(b).push_back(make_pair(a, c));
    }
    k--;
    dfs(k, k);
    for(int i = 0; i < q; i++){
        x = queries[i].first;
        y = queries[i].second;
        x--, y--;
        cout << pathTotal.at(x) + pathTotal.at(y) << endl;
    }
}
int main() {
    int n = 6, q = 3;
    k = 1;
    vector<tuple<int, int, int>> edges = {{1, 2, 2}, {1, 3, 4}, {3, 4, 2}, {3, 5, 3}, {5, 6, 2}};
    vector<pair<int, int>> queries = {{1, 4}, {2, 6}, {2, 5}};
    solve(n, q, edges, queries);
    return 0;
}

입력

6, 3, 1, {{1, 2, 2}, {1, 3, 4}, {3, 4, 2}, {3, 5, 3}, {5, 6, 2}}, {{1, 4}, {2, 6}, {2, 5}}

출력

6
11
9

복잡도 분석

인접 리스트 구축과 DFS 탐색에는 O(n)의 시간이 걸리며, 각 쿼리는 이미 계산된 거리 값을 더하기만 하므로 O(1)에 처리됩니다. 따라서 전체 시간 복잡도는 O(n + q)이고, 공간 복잡도는 인접 리스트와 거리 배열 저장에 O(n)입니다. 쿼리 개수가 많더라도 사전 계산 덕분에 매우 효율적으로 동작합니다.