Pythonで二分木における最大の完全部分木を見つける方法

問題の概要

二分木が与えられたとき、その木の中に含まれる最大の完全部分木(コンプリート・サブツリー)のサイズを求めることを考えます。

ここでいう完全二分木とは、最下層を除くすべてのレベルがノードで完全に埋め尽くされており、最下層のノードは可能な限り左側に配置されている二分木のことです。

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

このとき出力されるサイズは 4 となり、最大の完全部分木を通りがけ順(中順)で走査すると 10, 45, 60, 70, の順に出力されます。

解き方のアプローチ

この問題は、木を再帰的にたどりながら、各部分木が「完全(complete)」であるか「完全(perfect)」であるかを判定することで解けます。具体的には、次の手順に従います。

• 戻り値用の型として、フラグ isComplete と isPerfect(いずれも初期値は False)、さらに size(初期値 0)と rootTree(初期値 null)を定義します。
• ret_type := returnType()
• root が null の場合は、次のように設定して返します。
  • ret_type.isPerfect := True
  • ret_type.isComplete := True
  • ret_type.size := 0
  • ret_type.rootTree := None
• left_tree := checkCompleteness(root.left)
• right_tree := checkCompleteness(root.right)
• 左の部分木が perfect で、右の部分木が complete であり、両者の高さが等しい場合は、現在のノードを根とする木全体が complete になります。
  • ret_type.isComplete := True
  • ret_type.isPerfect := right_tree.isPerfect
  • ret_type.size := left_tree.size + right_tree.size + 1
  • ret_type.rootTree := root
• 左の部分木が complete で、右の部分木が perfect であり、左の高さが右よりちょうど1大きい場合も、現在のノードを根とする木は complete になります。
  • ret_type.isComplete := True
  • ret_type.isPerfect := False
  • ret_type.size := left_tree.size + right_tree.size + 1
  • ret_type.rootTree := root
• 上記のいずれの条件も満たさない場合は、現在のノードを根とする木は complete ではないため、左右の部分木のうちサイズが大きい方の結果をそのまま返します。
  • ret_type.isPerfect := False
  • ret_type.isComplete := False
  • ret_type.size := max(left_tree.size, right_tree.size)
  • left_tree.size > right_tree.size なら ret_type.rootTree := left_tree.rootTree、それ以外なら ret_type.rootTree := right_tree.rootTree

Pythonでの実装例

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

import math

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

class returnType:
    def __init__(self):
        self.isPerfect = None
        self.isComplete = None
        self.size = 0
        self.rootTree = None

def getHeight(size):
    return int(math.ceil(math.log(size + 1) / math.log(2)))

def checkCompleteness(root):
    ret_type = returnType()
    if root is None:
        ret_type.isPerfect = True
        ret_type.isComplete = True
        ret_type.size = 0
        ret_type.rootTree = None
        return ret_type
    left_tree = checkCompleteness(root.left)
    right_tree = checkCompleteness(root.right)
    if (left_tree.isPerfect and right_tree.isComplete
            and getHeight(left_tree.size) == getHeight(right_tree.size)):
        ret_type.isComplete = True
        ret_type.isPerfect = right_tree.isPerfect
        ret_type.size = left_tree.size + right_tree.size + 1
        ret_type.rootTree = root
        return ret_type
    if (left_tree.isComplete and right_tree.isPerfect
            and getHeight(left_tree.size) == getHeight(right_tree.size) + 1):
        ret_type.isComplete = True
        ret_type.isPerfect = False
        ret_type.size = left_tree.size + right_tree.size + 1
        ret_type.rootTree = root
        return ret_type
    ret_type.isPerfect = False
    ret_type.isComplete = False
    ret_type.size = max(left_tree.size, right_tree.size)
    if left_tree.size > right_tree.size:
        ret_type.rootTree = left_tree.rootTree
    else:
        ret_type.rootTree = right_tree.rootTree
    return ret_type

def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.data, end=', ')
        print_tree(root.right)

root = TreeNode(50)
root.left = TreeNode(30)
root.right = TreeNode(60)
root.left.left = TreeNode(5)
root.left.right = TreeNode(20)
root.right.left = TreeNode(45)
root.right.right = TreeNode(70)
root.right.left.left = TreeNode(10)

ans = checkCompleteness(root)
print("Size:", ans.size)
print("Inorder Traversal: ", end='')
print_tree(ans.rootTree)

入力

root = TreeNode(50)
root.left = TreeNode(30)
root.right = TreeNode(60)
root.left.left = TreeNode(5)
root.left.right = TreeNode(20)
root.right.left = TreeNode(45)
root.right.right = TreeNode(70)
root.right.left.left = TreeNode(10)

出力

Size: 4
Inorder Traversal: 10, 45, 60, 70,

このように、各ノードについて左右の部分木の状態を再帰的に評価していくことで、二分木の中から最大の完全部分木を効率よく見つけることができます。計算量は木のノード数を n とすると O(n) となり、非常に効率的なアプローチです。