【Python】二分探索木（BST）からK番目に小さい要素を求める方法

二分探索木（BST：Binary Search Tree）が与えられたとき、その木の中でK番目に小さい要素を見つける問題を考えてみましょう。

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

この木から3番目に小さい要素を求めたい場合は k = 3 となり、求める結果は 7 になります。

解法のアプローチ

この問題は、二分探索木の重要な性質を利用することで簡単に解けます。それは、BSTを中順走査（in-order traversal）すると、値が昇順に並んだリストが得られるというものです。

具体的には、以下の手順で解きます。

• 結果を格納するための空のリスト nodes を作成する
• solve(root, nodes) を呼び出して中順走査を行う
• 走査後の nodes の (k − 1) 番目の要素を返す

solve メソッドは、ルートノードとノード配列を受け取り、再帰的に次のように動作します。

• root が None（null）の場合は何もせず return する
• まず左部分木に対して solve(root.left, nodes) を再帰呼び出しする
• 現在のノード（root）の値を nodes リストに追加する
• 最後に右部分木に対して solve(root.right, nodes) を再帰呼び出しする

左 → 自身 → 右 の順で訪問するため、BSTの性質により値は必ず昇順でリストに格納されます。あとは k 番目の要素を取り出すだけです。

実装例

それでは、実際のPythonコードを見てみましょう。

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

def insert(temp, data):
    que = []
    que.append(temp)
    while (len(que)):
        temp = que[0]
        que.pop(0)
        if (not temp.left):
            temp.left = TreeNode(data)
            break
        else:
            que.append(temp.left)
        if (not temp.right):
            temp.right = TreeNode(data)
            break
        else:
            que.append(temp.right)

def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(Tree, element)
    return Tree

class Solution(object):
    def kthSmallest(self, root, k):
        nodes = []
        self.solve(root, nodes)
        return nodes[k-1]

    def solve(self, root, nodes):
        if root == None:
            return
        self.solve(root.left, nodes)
        nodes.append(root.data)
        self.solve(root.right, nodes)

ob1 = Solution()
tree = make_tree([10,5,15,2,7,13])
print(ob1.kthSmallest(tree, 3))

入力

[10,5,15,2,7,13]
3

出力

7

計算量について

このアプローチの計算量は以下の通りです。

• 時間計算量： O(n) — 木に含まれるすべてのノードを一度ずつ訪問するため
• 空間計算量： O(n) — 昇順の値をすべて保持するリストが必要なため（再帰のスタックも最大 O(h)、h は木の高さ）

補足：より効率的な方法

すべてのノードを訪問せずに済ませたい場合は、反復的な中順走査（スタックを使用）を行い、k 個目の要素に到達した時点で処理を打ち切る方法もあります。特に k が小さい場合や、木が非常に大きい場合には、この方法の方が効率的です。また、各ノードに「左部分木のノード数」を持たせておけば、O(h) の時間で K番目の要素を検索できるようになります。