【Python】二分探索木でk番目に小さい要素を効率的に求めるアルゴリズムと実装例

問題の概要

二分探索木（BST: Binary Search Tree）と整数 k が与えられたとき、木の中で k 番目に小さい値を見つけることを考えます。

例えば、次のような二分探索木があるとします。

        5
       / \
      4   10
         /  \
        7    15
       / \
      6   8

このとき k = 3 であれば、出力は 7 になります。

アプローチ：スタックを使った中順走査（In-order Traversal）

二分探索木には「中順走査を行うと、ノードを値の昇順に訪問できる」という重要な性質があります。この性質を利用すれば、「k 番目に小さい値を求める」という問題は「中順走査で k 番目に訪れたノードの値を返す」という問題に置き換えられます。

ここでは再帰呼び出しではなく、明示的なスタックを用いた反復処理で実装します。この方法なら、k 番目の要素に到達した時点ですぐに走査を打ち切れるため、木全体を巡回する無駄がありません。

アルゴリズムの手順

1. 空のスタックを用意し、カウンタ i := 0、答え ans := -1 で初期化する。
2. スタックが空でない、または root が null でない限り、以下を繰り返す。

• root が null でない間、root をスタックにプッシュし、root を左の子へ移動する（左端まで一気に辿る）。
• スタックから要素をポップし、v とする。
• i が k と一致していれば、ans に v の値を代入してループを抜ける。
• root を v の右の子に設定し、i を 1 増やす。

3. ans を返す。

※ この実装では k が 0 始まりのインデックスとして扱われます。つまり k = 0 が最小値、k = 3 は「4 番目に小さい値」（本例では 7）に対応します。1 始まりにしたい場合は、比較条件を i == k - 1 に変更してください。

Pythonによる実装例

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

class Solution:
    def solve(self, root, k):
        stack = []
        i = 0
        ans = -1
        while stack or root:
            while root:
                stack.append(root)
                root = root.left
            v = stack.pop()
            if i == k:
                ans = v.val
                break
            root = v.right
            i += 1
        return ans

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(10)
root.right.left = TreeNode(7)
root.right.right = TreeNode(15)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)
print(ob.solve(root, 3))

入力

root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(10)
root.right.left = TreeNode(7)
root.right.right = TreeNode(15)
root.right.left.left = TreeNode(6)
root.right.left.right = TreeNode(8)
k = 3

出力

7

計算量の評価

• 時間計算量：O(H + k)（H は木の高さ）。左端まで降りるのに O(H)、その後 k 回ポップするのに O(k) かかります。木が片側に偏っている最悪ケースでは O(n) になります。
• 空間計算量：O(H)。スタックに保持されるノード数は最大で木の高さ分です。

まとめ

二分探索木の「中順走査＝値の昇順アクセス」という性質を活かせば、k 番目に小さい要素を効率よく取り出せます。再帰実装も可能ですが、スタックを使った反復実装のほうが早期打ち切りが簡単で、深い木においてもスタックオーバーフローの心配が少なく安全です。