C++로 이진 탐색 트리(BST)에서 k번째로 작은 요소 찾기

문제 소개

이진 탐색 트리(Binary Search Tree, BST)와 정수 K가 입력으로 주어졌을 때, 트리에서 K번째로 작은 요소를 찾아야 합니다.

예를 들어 아래와 같은 BST가 있다고 가정해 보겠습니다.

k = 3일 때, 세 번째로 작은 값인 15가 출력됩니다.

접근 방법: 중위 순회 활용하기

BST의 가장 중요한 성질은 중위 순회(inorder traversal)를 수행하면 노드 값이 오름차순으로 방문된다는 점입니다. 이 성질을 활용하면 다음과 같은 알고리즘을 만들 수 있습니다.

find_kth_smallest() 함수를 정의합니다. 이 함수는 루트 노드(root), 방문한 노드 수(count), 목표 순번(k)을 매개변수로 받습니다.
루트가 NULL이면 NULL을 반환합니다.
왼쪽 서브트리를 먼저 재귀적으로 탐색합니다: left = find_kth_smallest(root->left, count, k)
left가 NULL이 아니라면 왼쪽 서브트리에서 이미 답을 찾은 것이므로 left를 그대로 반환합니다.
현재 노드를 방문했으므로 count를 1 증가시킵니다.
count가 k와 같다면 현재 노드가 바로 k번째로 작은 요소이므로 루트를 반환합니다.
아직 답을 찾지 못했다면 오른쪽 서브트리를 재귀적으로 탐색한 결과를 반환합니다.

메인 함수의 처리 흐름

count를 0으로 초기화합니다.
res = find_kth_smallest(root, count, k)를 호출합니다.
res가 NULL이면 "찾지 못함(Not found)"을 출력합니다.
그렇지 않으면 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

출력

동작 원리 살펴보기

위 예제 트리에서 중위 순회는 9 → 13 → 15 → 17 → 19 → 25 → 27 순서로 노드를 방문합니다. k = 3이므로 세 번째로 방문하는 노드인 15가 결과로 출력되는 것입니다. 답을 찾은 이후에는 재귀 호출 결과가 상위로 전파되며 더 이상의 탐색이 진행되지 않으므로 효율적입니다.

복잡도 분석

시간 복잡도: 최악의 경우 O(N)입니다. 다만 균형 잡힌 BST에서는 조기 종료 덕분에 O(H + K)(H는 트리의 높이)에 가깝게 동작합니다.
공간 복잡도: 재귀 호출 스택이 사용되므로 O(H)입니다.