파이썬으로 이진 탐색 트리(BST)에서 K번째로 작은 요소 찾기

문제 개요

이진 탐색 트리(Binary Search Tree, BST)가 하나 주어져 있다고 가정해 봅시다. 이때 트리에서 K번째로 작은 요소를 찾아야 하는 것이 목표입니다.

예를 들어 다음과 같은 트리가 있을 때 −

3번째로 작은 요소를 찾고 싶다면 k = 3이 되고, 정답은 7이 됩니다.

풀이 접근 방법

BST의 핵심 성질 중 하나는 중위 순회(inorder traversal)를 수행하면 노드의 값이 오름차순으로 정렬된 순서대로 방문된다는 점입니다. 이 성질을 활용하면 문제를 매우 간단하게 해결할 수 있습니다.

해결 절차는 다음과 같습니다 −

nodes라는 빈 리스트를 하나 생성합니다.
solve(root, nodes)를 호출하여 중위 순회를 진행합니다.
순회가 끝난 후 nodes 리스트의 (k − 1)번째 요소를 반환합니다.

solve 메서드의 동작 원리

solve 메서드는 루트 노드와 nodes 배열을 인자로 받으며, 다음과 같이 동작합니다 −

루트가 null이면 즉시 반환합니다(재귀 종료 조건).
왼쪽 서브트리에 대해 재귀 호출을 수행합니다.
현재 루트의 값을 nodes 리스트에 추가합니다.
오른쪽 서브트리에 대해 재귀 호출을 수행합니다.

중위 순회는 "왼쪽 → 루트 → 오른쪽" 순서로 노드를 방문하기 때문에, 순회가 완료되면 nodes 리스트에는 트리의 모든 값이 자동으로 오름차순으로 정렬되어 저장됩니다. 따라서 k번째로 작은 값은 곧 리스트의 (k − 1)번째 인덱스에 위치하게 됩니다.

구현 예제

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

출력

복잡도 분석

시간 복잡도: O(n) — 트리의 모든 노드를 한 번씩 방문해야 하므로 노드 수에 비례합니다.
공간 복잡도: O(n) — 모든 노드의 값을 저장하는 리스트와 재귀 호출 스택이 필요합니다.

만약 트리의 크기가 매우 크고 k가 작은 경우라면, 전체 순회 대신 중위 순회 도중 k번째 노드를 만나면 바로 종료하는 방식으로 최적화할 수도 있습니다.