Pythonで二分木の直径を求める方法【DFSを使った実装解説】

二分木の直径とは

二分木が与えられたとき、その木の直径(diameter)を計算することを考えます。二分木の直径とは、木の中の任意の2つのノードをつなぐ最長経路の長さのことです。

重要なポイントとして、この経路は必ずしも根(ルート)を通るとは限りません。

例えば、次のような木を考えてみましょう。

この場合、経路 [4, 2, 1, 3] または [5, 2, 1, 3] の長さが3本の辺で構成されているため、直径は3となります。

解法のアプローチ

この問題はDFS(深さ優先探索)を使うことで効率的に解くことができます。手順は以下の通りです。

• DFSで各ノードを訪問しながら直径を求めます。まず答えを格納する変数を answer = 0 で初期化します。
• ルートノードに対して dfs(root) を呼び出します。
• dfs(node) は以下のように動作します。
• ノードが存在しない場合は 0 を返します。
• left := 左部分木に対するDFSの結果、right := 右部分木に対するDFSの結果を取得します。
• answer := max(answer, left + right) として更新します。これは「現在のノードを頂点とした左右の経路の合計」が直径の候補になるためです。
• 親ノードに返す値は max(left + 1, right + 1) です。これは「そのノードから下方向へ伸びる最長の深さ」を表します。

Pythonでの実装例

それでは、実際のコードを見てみましょう。

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

def insert(temp, data):
    que = []
    que.append(temp)
    while len(que):
        temp = que[0]
        que.pop(0)
        if not temp.left:
            temp.left = TreeNode(data)
            break
        else:
            que.append(temp.left)
        if not temp.right:
            temp.right = TreeNode(data)
            break
        else:
            que.append(temp.right)

def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(Tree, element)
    return Tree

class Solution(object):
    def diameterOfBinaryTree(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        self.ans = 0
        self.dfs(root)
        return self.ans

    def dfs(self, node):
        if not node:
            return 0
        left = self.dfs(node.left)
        right = self.dfs(node.right)
        self.ans = max(self.ans, right + left)
        return max(left + 1, right + 1)

root = make_tree([1, 2, 3, 4, 5])
ob1 = Solution()
print(ob1.diameterOfBinaryTree(root))

入力

[1, 2, 3, 4, 5]

出力

3

計算量について

• 時間計算量: O(n) — 各ノードを一度だけ訪問するためです(n はノード数)。
• 空間計算量: O(h) — 再帰呼び出しのスタックの深さに依存します(h は木の高さ)。木が偏っている場合は最大でO(n)になります。

まとめ

二分木の直径を求める問題では、「各ノードを根とした左右の深さの合計」をDFSで順番に計算し、その最大値を記録していくのがポイントです。経路がルートを通らないケースにも対応できるよう、答えの更新を再帰処理の中で行っている点に注目してください。