q個のクエリで頂点kを経由する最短コストパスを求めるC++プログラム

問題概要

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 となります。

解法の考え方

グラフが最小連結、つまり木構造である場合、任意の2頂点間の経路は一意に決まります。頂点 k を経由するパスのコストは、「始点から k までの距離」と「k から目的地までの距離」を足し合わせたものになります。

そこで、頂点 k を起点として DFS（深さ優先探索）を一度だけ実行し、すべての頂点について k からの累積コストを前もって計算しておきます。こうすれば、各クエリには定数時間で回答できるようになります。

アルゴリズムの手順

n × n サイズのペアの2次元配列 graph を定義する
サイズ n の配列 pathTotal を定義する
関数 dfs(a, b) を定義する
    graph[a] 内の各要素 i について:
        i の first が b と等しい場合は何もせず次の反復へ進む
        pathTotal[i.first] := pathTotal[a] + i.second
        dfs(i.first, a) を呼び出す
i := 0 から n - 1 未満の間、i を 1 ずつ増やしながら繰り返す:
    a := edges[i] の第1要素
    b := edges[i] の第2要素
    c := edges[i] の第3要素
    a と b から 1 を引く
    graph[a] の末尾にペア (b, c) を追加する
    graph[b] の末尾にペア (a, c) を追加する
(k から 1 を引く)
dfs(k, k) を呼び出す
i := 0 から q 未満の間、i を 1 ずつ増やしながら繰り返す:
    x := queries[i] の第1要素
    y := queries[i] の第2要素
    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) となり、クエリ数が非常に多い場合でも高速に動作する効率的なアプローチです。