Pythonで二分木から最大の完全二分木（パーフェクトサブツリー）を見つける方法

与えられた二分木の中から、最大の完全二分木（Perfect Binary Tree）となっているサブツリーを見つける問題を考えてみましょう。完全二分木とは、すべての内部ノードが必ず2つの子を持ち、すべての葉ノードが同じ深さに位置する二分木のことです。

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

この場合の出力は 3 となり、見つかったサブツリーは次の通りです。

解法のアプローチ

この問題は、木を再帰的にたどりながら、各部分木について「完全二分木であるかどうか」と「高さ」を記録していくことで効率的に解けます。具体的な手順は以下の通りです。

• isPerfect（完全二分木かどうか）、height（高さ）、rootTree(サブツリーの根ノード)を保持するための構造体 RetType を定義します。初期値はすべて 0 です。
• 根ノード root を引数に取る関数 get_prefect_subtree() を定義します。
• r_type := 新しい RetType インスタンスを作成します。
• root が None の場合は、空の木は完全二分木とみなせるため、isPerfect を True、height を 0、rootTree を null に設定して返します。
• left_subtree := get_prefect_subtree(root.left) で左部分木を評価します。
• right_subtree := get_prefect_subtree(root.right) で右部分木を評価します。
• 左部分木と右部分木がどちらも完全であり、かつ両者の高さが等しい場合：
  • r_type.height := 左部分木の高さ + 1
  • r_type.isPerfect を True に設定
  • r_type.rootTree := 現在の root
  • r_type を返します。
• 上記の条件を満たさない場合は、r_type.isPerfect を False に設定します。
• r_type.height := 左右の部分木の高さの最大値を設定します。
• 左部分木の高さ > 右部分木の高さであれば、r_type.rootTree := left_subtree.rootTree とします。
• そうでなければ、r_type.rootTree := right_subtree.rootTree とします。
• r_type を返します。

ポイントは、左右の部分木がともに完全で高さが一致する場合のみ、現在のノードを根とする木全体が完全二分木になるという性質を利用している点です。最終的なサブツリーのノード数は、高さ h に対して 2^h − 1 で計算できます。

実装例

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

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.data = data
        self.left = left
        self.right = right
def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.data, end = ', ')
        print_tree(root.right)
class RetType:
    def __init__(self):
        isPerfect = 0
        height = 0
        rootTree = 0
def get_prefect_subtree(root):
    r_type = RetType()
    if (root == None) :
        r_type.isPerfect = True
        r_type.height = 0
        r_type.rootTree = None
        return r_type
    left_subtree = get_prefect_subtree(root.left)
    right_subtree = get_prefect_subtree(root.right)
    if (left_subtree.isPerfect and right_subtree.isPerfect and left_subtree.height == right_subtree.height) :
        r_type.height = left_subtree.height + 1
        r_type.isPerfect = True
        r_type.rootTree = root
        return r_type
    r_type.isPerfect = False
    r_type.height = max(left_subtree.height, right_subtree.height)
    if (left_subtree.height > right_subtree.height ):
        r_type.rootTree = left_subtree.rootTree
    else :
        r_type.rootTree = right_subtree.rootTree
    return r_type

root = TreeNode(2)
root.left = TreeNode(3)
root.right = TreeNode(4)
root.left.left = TreeNode(5)
root.left.right = TreeNode(6)
root.right.left = TreeNode(7)

res = get_prefect_subtree(root)
h = res.height

print ("Size: " , pow(2, h) - 1)
print ("Tree: ", end = " ")
print_tree(res.rootTree)

入力

root = TreeNode(2)
root.left = TreeNode(3)
root.right = TreeNode(4)
root.left.left = TreeNode(5)
root.left.right = TreeNode(6)
root.right.left = TreeNode(7)

出力

Size: 3
Tree: 5, 3, 6,

この結果から、ノード 5・3・6 から構成される高さ2の部分木が、この二分木における最大の完全二分木であることがわかります。計算量は木を一度だけ走査するため O(n) となり、大規模な木に対しても効率的に動作します。