C++で木構造における距離の総和を効率的に求める方法

N個のノードからなる無向かつ連結な木を考えます。ノードには0〜N-1のラベルが付いており、N-1本の辺が与えられています。i番目の辺は、ノードedges[i][0]とedges[i][1]を結んでいます。このとき、ans[i]が「ノードiとそれ以外のすべてのノードとの間の距離の合計」を表すようなリストansを求めるのがこの記事の目的です。

たとえば、入力がN = 6、edges = [(0,1),(0,2),(2,3),(2,4),(2,5)]である場合、出力は[8,12,6,10,10,10]となります。

解法のアプローチ

この問題を効率的に解くには、「再根付け(rerooting)」と呼ばれるテクニックを使います。まず最初のDFSで部分木ごとの情報(ノード数と距離の合計)を集め、次のDFSでその情報を流用することで、全ノードの答えを線形時間で求められます。

ステップ1:dfs1()の定義

引数としてnodeとparentを受け取る関数dfs1()を定義します。

• i := 0から始めて、i < graph[node]のサイズである限り、iを1ずつ増やしながら以下を繰り返します。
• child := graph[node][i]
• childがparentと異なる場合:
  • dfs1(child, node)を呼び出す
  • cnt[node] := cnt[node] + cnt[child]
  • ans[node] := ans[node] + cnt[child] + ans[child]

ステップ2:dfs2()の定義

同様に、nodeとparentを受け取る関数dfs2()を定義します。

• graph[node]の各要素childに対して、childがparentと異なる場合:
  • ans[child] := ans[node] - cnt[child] + N - cnt[child]
  • dfs2(child, node)を呼び出す

その他の準備

• 結果を格納する配列ansを定義する
• 部分木のノード数を格納する配列cntを定義する
• グラフの隣接リスト用に、10005行分の配列graphを定義する

メインメソッドでの処理の流れ

• メンバ変数Nに入力のNを格納する
• サイズNの配列ansを定義する
• サイズNの配列cntを定義し、すべての要素を1で初期化する
• n := edgesのサイズとする
• i := 0からn未満の間、以下を繰り返す:
  • u := edges[i][0]、v := edges[i][1]
  • graph[u]の末尾にvを挿入し、graph[v]の末尾にuを挿入する
• dfs1(0, -1)とdfs2(0, -1)を実行する
• ansを返す

アルゴリズムの仕組み

dfs1では、各ノードを根とする部分木に含まれるノード数(cnt)と、部分木内部だけでの距離の総和(ans)をボトムアップに計算します。続くdfs2では、親ノードですでに求めた答えを利用して子ノードの答えを導きます。根を隣接ノードへ移動させるとき、移動先の子部分木側に属するcnt[child]個のノードとの距離はそれぞれ1ずつ減り、それ以外のN - cnt[child]個のノードとの距離はそれぞれ1ずつ増えるため、ans[child] = ans[node] - cnt[child] + (N - cnt[child])という式で答えを更新できるのです。

C++実装例

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0;
   i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
class Solution {
   public:
   void dfs1(int node, int parent) {
      for (int i = 0; i < graph[node].size(); i++) {
         int child = graph[node][i];
         if (child != parent) {
            dfs1(child, node);
            cnt[node] += cnt[child];
            ans[node] += cnt[child] + ans[child];
         }
      }    
   }
   void dfs2(int node, int parent) {
      for (int i = 0; i < graph[node].size(); i++) {
         int child = graph[node][i];
         if (child != parent) {
            ans[child] = ans[node] - cnt[child] + N - cnt[child];
            dfs2(child, node);
         }
      }
   }
   vector<int> ans;
   vector<int> cnt;
   vector<int> graph[10005];
   int N;
   vector<int> sumOfDistancesInTree(int N, vector<vector<int> >& edges) {
      this->N = N;
      ans = vector<int>(N);
      cnt = vector<int>(N, 1);
      int n = edges.size();
      for (int i = 0; i < n; i++) {
         int u = edges[i][0];
         int v = edges[i][1];
         graph[u].push_back(v);
         graph[v].push_back(u);
      }
      dfs1(0, -1);
      dfs2(0, -1);
      return ans;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{0,1},{0,2},{2,3},{2,4},{2,5}}; print_vector(ob.sumOfDistancesInTree(6,    v));
}

入力

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

出力

[8, 12, 6, 10, 10, 10]