PythonでN分木の直径を求める方法｜DFSを使った実装例を解説

n分木（N-ary tree）が与えられ、その木の直径を求めることを考えます。木の直径とは、木に存在する任意の2つの葉ノードをつなぐ経路のうち、最も長いものを指します。この記事では、直径の長さを表す整数値を計算して返すプログラムをPythonで実装します。

問題の例

例えば、次のようなn分木が与えられた場合を考えてみましょう。

この場合の出力は 3 になります。

このn分木の直径は、27→14、14→42、そして42→56（または42→65）という辺から構成される経路です（図では赤線で示されています）。この経路の長さが3となるわけです。

解法のアプローチ

この問題は、深さ優先探索（DFS）の考え方を応用することで効率的に解けます。手順は以下の通りです。

1. 答えを保持する変数 ans を1で初期化します。
2. ノードの深さを再帰的に計算する関数 depth() を定義します。引数にはルートノードを受け取ります。
   • rootが空（None）の場合は0を返します。
   • 子ノードの深さを格納するリスト children を [0, 0] で初期化します。
   • 一時的なリスト temp_children を用意し、ルートの各子ノードに対して depth(child) の結果を順に追加していきます。
   • temp_children のサイズが0より大きければ、その内容を children に代入します。
   • children をソートし、末尾2つの要素（最も深い2つの子部分木）の合計に1を加えた値と ans を比較し、大きい方で ans を更新します。
   • children の最大値に1を加えた値を返します（これが現在のノードの深さになります）。
3. depth(root) を呼び出します。
4. 最後に ans - 1 を返します。

ここでのポイントは、あるノードを「経由」する最長経路は、そのノードの子部分木のうち深さが上位2つの部分木を通るときに最大化されるという点です。そのため、各ノードで上位2つの深さの合計を常に記録しておけば、木全体の直径を求められます。

Pythonでの実装例

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

class Node:
    def __init__(self, value, child = None) -> None:
        self.val = value
        self.children = []
        if child != None:
            for value in child:
                self.children.append(value)

ans = 1
def solve(root):
    def depth(root):
        global ans
        if not root:
            return 0
        children = [0, 0]
        temp_children = [depth(child) for child in root.children]
        if len(temp_children) > 0:
            children = temp_children
        ans = max(ans, sum(sorted(children)[-2:]) + 1)

        return max(children) + 1
    depth(root)

    return ans -1

node6 = Node(65)
node5 = Node(56)
node4 = Node(42, [node5, node6])
node3 = Node(32)
node2 = Node(27)
node1 = Node(14, [node2, node3, node4])
root = node1

print(solve(root))

入力

node6 = Node(65)
node5 = Node(56)
node4 = Node(42, [node5, node6])
node3 = Node(32)
node2 = Node(27)
node1 = Node(14, [node2, node3, node4])
root = node1

出力

3

まとめ

このように、DFSで各ノードの深さを計算しながら、途中で「上位2つの子部分木の深さの合計＋1」を追跡していくことで、n分木の直径を全ノードを一度ずつ訪問するだけの計算量（O(N)）で求められます。二分木だけでなく、任意の数の子を持つ木にもそのまま適用できる汎用的な手法なので、ぜひ覚えておきましょう。