Pythonで二分探索木（BST）の指定範囲内にあるノード数を求める方法

問題の概要

二分探索木（BST）が与えられ、さらに左側の境界値 l と右側の境界値 r が指定されます。このとき、木に含まれるすべてのノードの中で、値が l 以上 r 以下の範囲内にあるノードの個数を求めるのが目的です。

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

このとき l = 7、r = 13 とすると、範囲内に含まれるノードは 8、10、12 の3つなので、出力は 3 になります。

アルゴリズムの考え方

スタックを使った反復的な深さ優先探索（DFS）で木をたどります。重要なポイントは、二分探索木の性質を活かして枝刈り（pruning）を行うことです。値が境界より小さいノードの左側の部分木、あるいは境界より大きいノードの右側の部分木には、答えとなるノードが存在しないため、そこを探索せずにスキップすることで効率よくカウントできます。

• スタックを用意し、最初にルートをプッシュします。カウント用変数 count は 0 で初期化します。
• スタックが空になるまで、次の処理を繰り返します。
  • スタックの先頭要素を取り出し（ポップし）、node とします。
  • node が null でない場合、次のように分岐します。
    • ノードの値が l <= 値 <= r を満たす場合：count を 1 増やし、右の子と左の子の両方をスタックにプッシュします。
    • ノードの値が l より小さい場合：左側の部分木は範囲外になるため、右の子だけをスタックにプッシュします。
    • それ以外（ノードの値が r より大きい）の場合：右側の部分木は範囲外になるため、左の子だけをスタックにプッシュします。
• ループが終了したら count を返します。

Pythonでの実装例

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, l, r):
        stack, count = [root], 0
        while stack:
            node = stack.pop()
            if node:
                if l <= node.data <= r:
                    count += 1
                    stack += [node.right, node.left]
                elif node.data < l:
                    stack += [node.right]
                else: stack += [node.left]
        return count
ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root, 7,13))

計算量について

最悪の場合（木が大きく偏っている、または範囲が非常に広い場合）はすべてのノードを訪問するため、時間計算量は O(n) となります。ただし、範囲が狭い場合は枝刈りにより訪問するノード数が大幅に減少し、実際の探索コストは小さくなります。空間計算量はスタックの深さに依存し、バランスの取れた木では O(log n)、最悪ケースでは O(n) です。

入力例

root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
7,13

出力

3