Pythonで木構造内の距離がちょうどkとなる頂点ペアの個数を求める方法

問題の概要

整数 k と、n 個のノードからなる木(ツリー構造)が与えられたとします。このとき、頂点間の距離がちょうど k となる異なる頂点ペアの総数を数えるのが課題です。

例として、k = 2 の場合に次のような木が与えられたとします。

このとき、出力は 4 になります。実際、距離がちょうど 2 となるペアは (1, 3)、(1, 5)、(2, 4)、(3, 5) の 4 組です。

解法の考え方

この問題は、深さ優先探索(DFS)を用いたボトムアップの集計によって効率よく解けます。ポイントは、各頂点ごとに「その頂点から距離 j にある子孫ノードの個数」を記録しておき、部分木を統合しながら答えを加算していくことです。

アルゴリズムで使用する主な変数は次の通りです。

• N = 5005:ノード数の上限
• graph:サイズ N の隣接リスト
• vertex_count:505 × 5005 の二次元配列。vertex_count[v][j] は「頂点 v の部分木内で、v から距離 j の位置にあるノードの数」を表す
• res = 0:最終的な答えを格納するカウンタ

辺の追加:insert_edge() 関数

頂点 x と y を結ぶ無向辺をグラフに登録します。具体的には、graph[x] の末尾に y を、graph[y] の末尾に x をそれぞれ追加します。

DFSによる集計:dfs() 関数

現在の頂点 v とその親 parent を引数にとり、再帰的に処理を進めます。

1. vertex_count[v][0] = 1 と設定する(自分自身との距離は 0 であるため)。
2. v に隣接する各頂点 i について、i が親でない場合は次を実行する。
   • dfs(i, v) を呼び出し、子 i の部分木の情報を先に完成させる。
   • j を 1 から k まで動かしながら res += vertex_count[i][j-1] * vertex_count[v][k-j] を計算する。これは「これまでに処理済みの子孫」と「新しくマージする部分木 i」をつなぐペアの数を数える操作。
   • 続いて j を 1 から k まで動かし、vertex_count[v][j] += vertex_count[i][j-1] により、子 i の距離情報を頂点 v に統合する。

このように、各部分木を順番にマージしながらペアを数えることで、同じペアを二重にカウントすることなく正確な答えが得られます。計算量は O(n・k) となり、比較的大規模な木でも高速に動作します。

実装例(Python)

N = 5005
graph = [[] for i in range(N)]
vertex_count = [[0 for i in range(505)] for i in range(N)]
res = 0

def insert_edge(x, y):
    graph[x].append(y)
    graph[y].append(x)

def dfs(v, parent):
    global res
    vertex_count[v][0] = 1
    for i in graph[v]:
        if (i != parent):
            dfs(i, v)
            for j in range(1, k + 1):
                res += vertex_count[i][j - 1] * vertex_count[v][k - j]
            for j in range(1, k + 1):
                vertex_count[v][j] += vertex_count[i][j - 1]

k = 2
insert_edge(1, 2)
insert_edge(2, 3)
insert_edge(3, 4)
insert_edge(2, 5)
dfs(1, 0)
print(res)

入力

k = 2
insert_edge(1, 2)
insert_edge(2, 3)
insert_edge(3, 4)
insert_edge(2, 5)

出力

4

まとめ

木構造上で距離がちょうど k となる頂点ペアを数える問題は、DFS で各頂点からの距離別ノード数を管理しながら部分木を統合する手法が有効です。本記事のコードでは、vertex_count 配列を活用することで O(n・k) の計算量で解けることを確認しました。競技プログラミングでも頻出のテクニックなので、ぜひ理解を深めておきましょう。