Pythonで木の特定の辺を含む一意なパスの総数をカウントするプログラム

木構造を表す辺のリスト (u, v) が与えられます。ここで、各辺について「その辺を含む一意なパス(単純パス)」の総数を求め、入力された辺と同じ順序で結果を返す必要があります。

例として、入力が edges = [[0, 1], [0, 2], [1, 3], [1, 4]] の場合を考えてみましょう。

この場合、出力は [6, 4, 4, 4] となります。

解き方のアプローチ

この問題は、以下の手順で解くことができます。

• 与えられた辺から隣接リスト adj を作成します。

• 各頂点の部分木サイズを記録するためのマップ count を用意します。

• 関数 dfs(x, parent) を定義します。

• count[x] を 1 で初期化します。

• adj[x] 内の各隣接頂点 nb について、次を繰り返します。

  • nb が parent と同じ場合はスキップします。

  • それ以外の場合、count[x] に dfs(nb, x) の戻り値を加算します。

• count[x] を返します。

• メイン処理では、まず dfs(0, -1) を呼び出して全頂点の部分木サイズを計算します。

• 結果格納用の空リスト ans を作成します。

• edges 内の各辺 (a, b) について、次を行います。

  • x := count[a] と count[b] の小さい方

  • ans の末尾に x * (count[0] - x) を追加します。

• ans を返します。

なぜこの方法で求まるのか

木から1本の辺を取り除くと、木は2つの連結成分に分割されます。ある辺を含むパスは必ず「一方の成分内の頂点」から始まり「もう一方の成分内の頂点」で終わるため、その辺を通るパスの数は「片側の頂点数 × 反対側の頂点数」として計算できます。

DFSによって根(頂点0)を基準とした部分木サイズを事前に求めておけば、辺 (a, b) における小さい側のサイズ x は min(count[a], count[b]) で得られ、全体の頂点数 n = count[0] であることから、答えは x * (n − x) と表せます。計算量は O(N) と非常に効率的です。

実装例

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

from collections import defaultdict
class Solution:
    def solve(self, edges):
        adj = defaultdict(list)
        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)
        count = defaultdict(int)
        def dfs(x, parent):
            count[x] = 1
            for nb in adj[x]:
                if nb == parent:
                    continue
                count[x] += dfs(nb, x)
            return count[x]
        dfs(0, -1)
        ans = []
        for a, b in edges:
            x = min(count[a], count[b])
            ans.append(x * (count[0] - x))
        return ans
ob = Solution()
edges = [
    [0, 1],
    [0, 2],
    [1, 3],
    [1, 4]
]
print(ob.solve(edges))

入力

[
    [0, 1],
    [0, 2],
    [1, 3],
    [1, 4]
]

出力

[6, 4, 4, 4]