Pythonで二分木の葉ノードと非葉ノードの数を求めるプログラム

二分木が与えられたとき、最初の要素に葉ノード（リーフノード）の数、2番目の要素に非葉ノードの数を格納した2つの数値のペアを求める問題を考えてみましょう。

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

この木には葉ノードが3つ、非葉ノードが2つ存在するため、出力は (3, 2) となります。

解き方のアルゴリズム

この問題は、再帰処理を使って以下の手順で解くことができます。

• ノード n が null（None）である場合は、(0, 0) を返します。
• n の左の子と右の子がどちらも null の場合（つまり n が葉ノードの場合）は、(1, 0) を返します。
• left := solve(n の左の子)
• right := solve(n の右の子)
• (left[0] + right[0], 1 + left[1] + right[1]) を返します。第1要素は左右の部分木から得られる葉ノード数の合計、第2要素は自分自身を含む非葉ノードの数になります。

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

実装例

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.val = data
        self.left = left
        self.right = right
class Solution:
    def solve(self, n):
        if not n:
            return 0, 0
        if not n.left and not n.right:
            return 1, 0
        left, right = self.solve(n.left), self.solve(n.right)
        return left[0] + right[0], 1 + left[1] + right[1]
ob = Solution()
root = TreeNode(6)
root.left = TreeNode(2)
root.right = TreeNode(6)
root.right.left = TreeNode(10)
root.right.right = TreeNode(2)
print(ob.solve(root))

入力

root = TreeNode(6)
root.left = TreeNode(2)
root.right = TreeNode(6)
root.right.left = TreeNode(10)
root.right.right = TreeNode(2)

出力

(3, 2)