C++で重み付きグラフにおける真のクエリ数をカウントするプログラム

問題の概要

無向グラフの辺リストを考えます。各辺は [u, v, w] という形式で与えられ、u と v はそれぞれ始点・終点の頂点、w はその辺の重みを表します。さらに、同じ形式 [u, v, w] のクエリリストも与えられます。各クエリは「頂点 u から v へ、パス上のすべての辺の重みが w 以下であるような経路は存在するか?」という問いを意味します。このとき、真(true)となるクエリの個数を求めてください。

たとえば、入力が以下の場合を考えてみましょう。

edges = [[0, 1, 6], [1, 2, 7], [2, 3, 8], [0, 3, 5]]
queries = [[0, 2, 14], [1, 0, 3]]

この場合の出力は 1 になります。クエリ [0, 2, 14] については、経路 [0 → 1 → 2] をたどることで頂点 0 から 2 へ移動できます。この経路の辺の重みは 6 と 7(合計 13)であり、どちらも上限 14 以下なので成立します。一方、クエリ [1, 0, 3] については、重み 3 以下の辺だけを使って頂点 1 から 0 へ到達する経路は存在しないため、偽(false)となります。

解法のアプローチ

この問題は、Union-Find(素集合データ構造 / DSU)を使った「オフライン処理」によって効率的に解くことができます。基本となるアイデアは、辺とクエリをひとつのリストにまとめて重みの昇順にソートし、重みの小さいものから順に処理していくというものです(Kruskal 法による最小全域木の構築と同じ発想です)。こうすることで、あるクエリを処理する時点では、それ以前に追加された辺(重みがクエリの閾値以下の辺)だけがグラフに含まれていることが保証されます。したがって、クエリの両端点が同じ連結成分に属していれば、そのクエリは真であると判定できます。

アルゴリズムの手順

• 関数 get_parent() を定義します。引数は頂点 x と親管理用の配列 par です。
  • par[x] が x 自身でない場合は、par[x] := get_parent(par[x], par) として再帰的に根を辿り、経路圧縮を行います。
• return par[x]
• メインメソッドでは以下を実行します。
• 2 次元配列 gr を定義します。
• n := 0 と初期化します。
• edges 内の各辺 t について:
  • n := max(n, t[0], t[1]) として頂点番号の最大値を更新します。
  • gr に行 [t[2], 0, t[0], t[1]] を挿入します(第1要素が重み、第2要素の 0 は「辺」を示すフラグ)。
• queries 内の各クエリ t について:
  • gr に行 [t[2], 1, t[0], t[1]] を挿入します(第2要素の 1 は「クエリ」を示すフラグ)。
• gr をソートします(重み順に並び、同じ重みの場合はフラグの小さい「辺」が先に処理されます)。
• サイズ n + 1 の配列 par を定義し、-1 で初期化します。
• i := 0 から i <= n まで繰り返し:
  • par[i] := i と設定します(初期状態では自分自身が親)。
• sz := クエリの件数、ans := 0 とします。
• gr の各行 t に対して:
  • a := t[2]、b := t[3]、tp := t[1]、d := t[0] を取り出します。
  • px := get_parent(a, par)、py := get_parent(b, par) でそれぞれの根を取得します。
  • tp が 0(辺)の場合:
    • px ≠ py であれば、par[py] := px として 2 つの連結成分を統合します。
  • それ以外(クエリ)の場合:
    • px == py であれば、両端点がすでに接続されているため ans を 1 増やします。
  • sz を 1 減らし、sz が 0 になった時点でループを抜けます(すべてのクエリを処理し終えたための早期終了)。
• return ans

計算量

辺の数を E、クエリの数を Q とすると、M = E + Q 個の要素のソートに O(M log M)、その後の処理は Union-Find のほぼ定数時間の操作により O(M α(M)) で完了します。クエリごとに毎回グラフを探索する方法と比べて大幅に高速なのがポイントです。

C++ 実装例

理解を深めるために、以下の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
int get_parent(int x, vector<int>& par) {
   if (par[x] != x) {
      par[x] = get_parent(par[x], par);
   }
   return par[x];
}
int solve(vector<vector<int>>& edges, vector<vector<int>>& queries) {
   vector<vector<int>> gr;
   int n = 0;
   for (auto t : edges) {
      n = max(n, max(t[0], t[1]));
      gr.push_back({t[2], 0, t[0], t[1]});
   }
   for (auto t : queries) {
      gr.push_back({t[2], 1, t[0], t[1]});
   }
   sort(gr.begin(), gr.end());
   vector<int> par(n + 1, -1);
   for (int i = 0; i <= n; i++) {
      par[i] = i;
   }
   int sz = queries.size();
   int ans = 0;
   for (auto t : gr) {
      int a = t[2];
      int b = t[3];
      int tp = t[1];
      int d = t[0];
      int px = get_parent(a, par);
      int py = get_parent(b, par);
      if (tp == 0) {
         if (px != py) {
            par[py] = px;
         }
      }else {
         if (px == py) {
            ans++;
         }
         sz--;
         if(sz == 0) {
            break;
         }
      }
   }
   return ans;
}
int main(){
   vector<vector<int>> edges = {{0, 1, 6},{1, 2, 7},{2, 3, 8},{0, 3, 5}};
   vector<vector<int>> queries = {{0, 2, 14},{1, 0, 3}};
   cout << solve(edges, queries);
}

入力

{{0, 1, 6},{1, 2, 7},{2, 3, 8},{0, 3, 5}}, {{0, 2, 14},{1, 0, 3}}

出力

1