Pythonで二分木に含まれる一人っ子ノードの数を数えるプログラム

二分木が与えられたとき、「一人っ子ノード」——つまり、親が持つ唯一の子であるようなノード——の数を求めることを考えます。あるノード x が一人っ子ノードであるとは、その親ノードが x だけを子として持ち、左か右のどちらか片方しか子を持たない場合を指します。

例として、次のような二分木を考えてみます。

この木では、ノード 8 とノード 6 がそれぞれ親(7 と 10)の唯一の子となっているため、出力は 2 になります。

アルゴリズムの流れ

幅優先探索(BFS)を用いて、木を上から順に走査しながら一人っ子ノードを数えていきます。手順は以下の通りです。

• ルートが null の場合は 0 を返す
• 両端キュー(deque)を作成し、末尾にルートを挿入する
• カウント用変数 count を 0 で初期化する
• キューが空になるまで、次の処理を繰り返す
  • キューの先頭からノードを取り出し、削除する
  • current の左の子が存在する場合は、その左の子をキューに挿入し、右の子が存在しなければ count を 1 増やす
  • current の右の子が存在する場合は、その右の子をキューに挿入し、左の子が存在しなければ count を 1 増やす

最後に count を返します。それでは、実際の実装を見てみましょう。

サンプルコード

from collections import deque

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):
        if not root:
            return 0
        d = deque()
        d.append(root)
        count = 0
        while d:
            current = d.popleft()
            if current.left:
                d.append(current.left)
                if not current.right:
                    count += 1
            if current.right:
                d.append(current.right)
                if not current.left:
                    count += 1
        return count

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

入力

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

出力

2

計算量について

このアルゴリズムでは、各ノードをちょうど一度ずつ訪問するため、時間計算量は O(n)(n はノードの総数)です。また、最悪の場合(左右に大きく偏った木など)にはキューに多数のノードが格納される可能性があるため、空間計算量も O(n) となります。