Pythonで「ほぼBST」を正確な二分探索木（BST）へ修正するプログラム

ここでは、2つのノードの値だけが入れ替わってしまった二分木（ほぼBST）を、正しい二分探索木（BST）に復元する方法を解説します。BSTでは中順走査（inorder traversal）を行うと値が必ず昇順に並ぶという性質があるため、この性質を利用して入れ替わったノードを検出・修正できます。

例えば、次のような入力が与えられたとします。

これを修正すると、出力は次のようになります。

解決のための手順

この問題は、以下のアルゴリズムで解決できます。

• 変数を初期化する：prev_node := null、min_node := null、max_node := null
• found_one := False とする
• root の中順走査で各ノードを順に処理する：
  • prev_node が null でない場合：
    • node の値 < prev_node の値（順序違反）の場合：
      • min_node が null、または node の値 < min_node の値なら、min_node := node
      • max_node が null、または max_node の値 < prev_node の値なら、max_node := prev_node
      • found_one が True なら、ループを抜ける
      • そうでなければ、found_one := True とする
  • prev_node := node を更新する
• min_node と max_node の値を入れ替える
• root を返す

ポイントは、中順走査中に「現在のノードの値が直前のノードの値より小さい」という順序違反を最大2回まで検出することです。1回目の違反で、大きい側のノード（max_node）と小さい側のノード（min_node）の候補を記録し、2回目の違反が見つかった時点で走査を終了します。最後に両者の値を交換すれば、元のBSTが復元されます。計算量はO(n)、空間計算量も再帰スタックを除けばO(1)と効率的です。

それでは、以下の実装を見て理解を深めましょう。

実装例

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.val = data
        self.left = left
        self.right = right
    
def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.val, end = ', ')
        print_tree(root.right)
    
def __iter__(self):
    if self.left:
        for node in self.left:
            yield node
    yield self
    if self.right:
        for node in self.right:
            yield node

setattr(TreeNode, "__iter__", __iter__)
class Solution:
    def solve(self, root):
        prev_node = None
        min_node = None
        max_node = None
        found_one = False
        for node in root:
            if prev_node:
                if node.val < prev_node.val:
                    if min_node is None or node.val < min_node.val:
                        min_node = node
                    if max_node is None or max_node.val < prev_node.val:
                        max_node = prev_node
                    if found_one:
                        break
                    else:
                        found_one = True
            prev_node = node
        min_node.val, max_node.val = max_node.val, min_node.val
        return root
    
ob = Solution()
root = TreeNode(3)
root.left = TreeNode(6)
root.right = TreeNode(8)
root.right.left = TreeNode(2)
root.right.right = TreeNode(9)
print_tree(ob.solve(root))

入力

root = TreeNode(3)
root.left = TreeNode(6)
root.right = TreeNode(8)
root.right.left = TreeNode(2)
root.right.right = TreeNode(9)

出力

2, 3, 6, 8, 9,

このように、ジェネレータによる中順走査で順序違反の箇所を特定し、入れ替わった2つのノードの値を交換するだけで、壊れたBSTを簡単かつ効率的に修復できることが分かります。