Pythonで二分探索木（BST）から指定範囲外のノードをすべて削除する方法

問題の概要

二分探索木（BST）と2つの値 low、high が与えられたとき、[low, high] の範囲（境界値を含む）に該当しないノードをすべて木から削除するプログラムを作成します。

例として、次のようなBSTを考えてみましょう。

ここで low = 7、high = 10 とした場合、範囲外のノード（5 や 1 など）が削除され、出力は次のようになります。

解法のアプローチ

この問題は再帰を利用することで簡潔に解くことができます。手順は以下の通りです。

• 関数 solve() を定義します。引数は root（現在のノード）、low、high の3つです。
• root が null（空）の場合は何もせずに返します。
• low が現在のノードの値より大きい場合、このノードとその左部分木は範囲外となるため、右部分木に対して solve() を再帰的に呼び出した結果を返します。
• high が現在のノードの値より小さい場合、このノードとその右部分木は範囲外となるため、左部分木に対して solve() を再帰的に呼び出した結果を返します。
• 上記以外の場合（ノードが範囲内の場合）は、右部分木と左部分木それぞれに対して solve() を呼び出し、結果を各子ポインタに代入してから root を返します。

実装例

それでは、実際の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 Solution:
    def solve(self, root, low, high):
        if not root:
            return
        if low > root.data:
            return self.solve(root.right, low, high)
        if high < root.data:
            return self.solve(root.left, low, high)
        root.right = self.solve(root.right, low, high)
        root.left = self.solve(root.left, low, high)
        return root

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(1)
root.right = TreeNode(9)
root.right.left = TreeNode(7)
root.right.right = TreeNode(10)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)
low = 7
high = 10
ret = ob.solve(root, low, high)
print_tree(ret)

入力

root = TreeNode(5)
root.left = TreeNode(1)
root.right = TreeNode(9)
root.right.left = TreeNode(7)
root.right.right = TreeNode(10)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)
low = 7
high = 10

出力

7, 8, 9, 10,

処理の流れの解説

このアルゴリズムでは、BSTの性質（左の子 < 親 < 右の子）を活かしています。あるノードの値が low 未満であれば、その左部分木全体も必ず low 未満になるため、まとめて切り捨てることができます。同様に、ノードの値が high より大きければ、右部分木全体を切り捨てられます。これにより、不要な探索を省きながら効率的に範囲外のノードを除去できます。