Pythonで二分木の各レベルの最大幅を求めるプログラム

二分木が与えられたとき、ツリー内の任意のレベルにおける最大幅を求めることを考えます。ここでいう「レベルの幅」とは、そのレベルにおいて最も左端にあるノードと最も右端にあるノードの間に含まれるノード数のことです。

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

この場合、出力は 2 となります。

解決のための手順

この問題を解くために、以下の手順に従います。

• 各深さにおける位置の最小値と最大値を保持するマップ d を作成します。初期値は、最小値を無限大（∞）、最大値を 0 とします。
• 関数 dfs() を定義します。この関数は引数として root、pos := 0、depth := 0 を受け取ります。
• root が null の場合は何もせずに return します。
• d[depth][0] を d[depth][0] と pos の最小値で更新します。
• d[depth][1] を d[depth][1] と pos の最大値で更新します。
• 左の子ノードに対して dfs(node.left, 2*pos, depth+1) を呼び出します。
• 右の子ノードに対して dfs(node.right, 2*pos+1, depth+1) を呼び出します。

メインメソッドでの処理

• dfs(root) を呼び出して探索を開始します。
• 変数 mx := 0 を初期化します。
• マップ d のすべての値（最小値・最大値のペア）について以下を繰り返します。
  • left := min、right := max とします。
  • mx を max(mx, right - left + 1) で更新します。
• 最後に mx を返します。

それでは、実際の実装を見て理解を深めましょう。

実装例

from collections import defaultdict

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

class Solution:
    def solve(self, root):
        d = defaultdict(lambda: [1e9, 0])

        def dfs(node, pos=0, depth=0):
            if not node:
                return
            d[depth][0] = min(d[depth][0], pos)
            d[depth][1] = max(d[depth][1], pos)
            dfs(node.left, 2 * pos, depth + 1)
            dfs(node.right, 2 * pos + 1, depth + 1)

        dfs(root)
        mx = 0
        for interval in d.values():
            l, r = interval
            mx = max(mx, r - l + 1)
        return mx

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(1)
root.right = TreeNode(9)
root.right.left = TreeNode(7)
root.right.right = TreeNode(10)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)
print(ob.solve(root))

入力

root = TreeNode(5)
root.left = TreeNode(1)
root.right = TreeNode(9)
root.right.left = TreeNode(7)
root.right.right = TreeNode(10)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)

出力

2

アルゴリズムのポイント

このアプローチでは、完全二分木のように各ノードにインデックス（左の子は 2*pos、右の子は 2*pos+1）を割り当てることで、同じ深さにあるノード同士の相対的な位置関係を数値で表現できます。これにより、各深さにおける最左ノードと最右ノードの位置差から、そのレベルの幅を正確に計算することが可能になります。