【C++】二分探索木（BST）でk番目に小さい要素を検索する方法

問題概要

二分探索木（BST）と整数 k が入力として与えられたとき、木の中で k番目に小さい要素 を見つける問題を解説します。

例えば、以下のようなBSTを考えてみましょう。

この木に対して k = 3 を指定した場合、出力は 15 になります。木の要素を昇順に並べると「9, 13, 15, 17, 19, 25, 27」となり、3番目の値が15であるためです。

アルゴリズムの考え方

二分探索木には、「中順走査（in-order traversal）」を行うと要素が昇順に訪問されるという重要な性質があります。この性質を利用し、走査中に訪問したノード数をカウントしていき、k番目に到達した時点でそのノードを返せば、目的の要素を効率的に求められます。

具体的な手順は以下の通りです。

• 関数 find_kth_smallest() を定義する。引数は root（現在のノード）、count（参照渡しのカウンタ）、k の3つです。

• root が NULL の場合は NULL を返します。

• まず左部分木に対して find_kth_smallest() を再帰的に呼び出します。

• 左側から結果（NULL以外）が返ってきた場合は、それをそのまま返します。

• カウンタ count を1増やします。

• count が k と等しくなったら、現在のノード（root）を返します。

• そうでなければ、右部分木に対して再帰的に呼び出した結果を返します。

メイン関数での処理

• カウンタ count を 0 で初期化します。

• res = find_kth_smallest(root, count, k) を実行します。

• res が NULL の場合は「見つからない」ことを表示します。

• それ以外の場合は、res の値を表示します。

C++による実装例

それでは、実際のコードを見て理解を深めましょう。

#include <iostream>
using namespace std;
struct TreeNode {
    int val;
    TreeNode *left, *right;
    TreeNode(int x) {
        val = x;
        left = right = NULL;
    }
};
TreeNode* find_kth_smallest(TreeNode* root, int &count, int k) {
    if (root == NULL)
        return NULL;
    TreeNode* left = find_kth_smallest(root->left, count, k);
    if (left != NULL)
        return left;
    count++;
    if (count == k)
        return root;
    return find_kth_smallest(root->right, count, k);
}
void kth_smallest(TreeNode* root, int k) {
    int count = 0;
    TreeNode* res = find_kth_smallest(root, count, k);
    if (res == NULL)
        cout << "Not found";
    else
        cout << res->val;
}
int main() {
    TreeNode* root = new TreeNode(25);
    root->left = new TreeNode(13);
    root->right = new TreeNode(27);
    root->left->left = new TreeNode(9);
    root->left->right = new TreeNode(17);
    root->left->right->left = new TreeNode(15);
    root->left->right->right = new TreeNode(19);

    int k = 3;
    kth_smallest(root, k);
}

入力

TreeNode* root = new TreeNode(25); root->left = new TreeNode(13);
root->right = new TreeNode(27); root->left->left = new
TreeNode(9); root->left->right = new TreeNode(17); root->left->right->left = new TreeNode(15); root->left->right->right = new TreeNode(19); k = 3

出力

15

計算量について

このアルゴリズムの時間計算量は O(k) です。k番目のノードが見つかった時点で再帰が早期に終了するため、木全体を走査する必要がありません。また、空間計算量は木の高さに依存し、バランスの取れたBSTでは O(log n)、極端に偏った木では最悪 O(n) となります。