Pythonで二分木のすべての葉が同じレベルにあるかどうかを確認するプログラム

二分木が与えられたとき、そのすべての葉（リーフノード）が同じレベル（同じ深さ）に存在するかどうかを判定する問題を考えてみましょう。

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

出力は True になります。

解決のアプローチ

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

• dfs() という関数を定義します。この関数は root（現在のノード）と d（現在の深さ）を受け取ります。
• root が null でない場合、以下の処理を行います。
  • root の左の子と右の子がどちらも null の場合（つまり葉ノードの場合）、d を depth リストの末尾に追加します。
  • それ以外の場合は、再帰的に探索を続けます。
    • dfs(root の左の子, d + 1)
    • dfs(root の右の子, d + 1)
• メインの処理では、以下を実行します。
• depth := 新しい空のリストを作成
• dfs(root, 0) を呼び出して全ノードを走査
• depth リストに含まれる値が1種類だけなら True を返す（すべての葉が同じ深さにあることを意味する）

実装例

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

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

class Solution:
    def solve(self, root):
        self.depth = []
        self.dfs(root, 0)
        return len(set(self.depth)) == 1

    def dfs(self, root, depth):
        if root:
            if not root.left and not root.right:
                self.depth.append(depth)
            else:
                self.dfs(root.left, depth + 1)
                self.dfs(root.right, depth + 1)

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(4)
root.left.left = TreeNode(2)
root.right = TreeNode(10)
root.right.left = TreeNode(7)
root.right.right = TreeNode(15)
print(ob.solve(root))

入力

root = TreeNode(5)
root.left = TreeNode(4)
root.left.left = TreeNode(2)
root.right = TreeNode(10)
root.right.left = TreeNode(7)
root.right.right = TreeNode(15)

出力

True

コードのポイント

このアルゴリズムの核心は、DFS（深さ優先探索）を使って木を走査し、葉ノードに到達した時点での深さを記録することです。最後に set() を使って depth リスト内の値を重複排除し、要素数が1であれば「すべての葉が同じ深さにある」と判断できます。計算量は木のノード数に対して O(n)、空間計算量も O(n) となり、効率的な解法です。