Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 이진 탐색 트리(BST)에서 k번째로 작은 요소 찾는 방법

이진 탐색 트리(Binary Search Tree)와 정수 k가 주어졌을 때, 트리 안에서 k번째로 작은 값을 찾아야 하는 문제입니다.

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

Python으로 이진 탐색 트리(BST)에서 k번째로 작은 요소 찾는 방법

이때 k = 3이라면, 세 번째로 작은 값인 7이 출력됩니다.

문제 해결 접근 방식

이 문제는 중위 순회(In-order Traversal)를 활용하면 효율적으로 해결할 수 있습니다. 이진 탐색 트리를 중위 순회하면 노드 값이 오름차순으로 방문되기 때문에, 순회 도중 k번째로 방문한 노드의 값이 곧 k번째로 작은 값이 됩니다.

스택을 사용한 반복적(iterative) 중위 순회로 구현하는 절차는 다음과 같습니다.

  • 빈 스택(stack)을 준비합니다.

  • 카운터 i := 0, 결과값 ans := -1로 초기화합니다.

  • 스택이 비어 있지 않거나 root가 null이 아닌 동안 다음을 반복합니다.

    • root가 null이 아닌 동안 root를 스택에 push하고, root를 왼쪽 자식으로 이동시킵니다.

    • 스택에서 요소 v를 pop 합니다.

    • i가 k와 같다면, ans에 v의 값을 저장하고 루프를 종료합니다.

    • 그렇지 않으면 root를 v의 오른쪽 자식으로 이동시키고, i를 1 증가시킵니다.

  • 반복이 끝나면 ans를 반환합니다.

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)
3

출력

7

코드 설명 및 시간 복잡도

위 코드는 스택을 이용해 왼쪽 서브트리부터 차례대로 탐색합니다. 각 노드를 pop 할 때마다 카운터 i를 증가시키며, i가 k에 도달하는 순간 해당 노드의 값을 결과로 저장하고 즉시 종료합니다. 덕분에 트리 전체를 순회할 필요 없이 k번째 노드까지만 탐색하면 됩니다.

  • 시간 복잡도: O(H) — H는 트리의 높이입니다. 최악의 경우에도 전체 노드 수 N보다는 작거나 같습니다.
  • 공간 복잡도: O(H) — 스택에 최대 트리 높이만큼의 노드가 저장됩니다.

만약 재귀(recursion) 방식을 선호한다면, 중위 순회 함수를 재귀적으로 호출하면서 방문 순서를 세는 방법으로도 동일하게 구현할 수 있습니다. 다만 트리가 매우 깊은 경우에는 위와 같은 반복 방식이 스택 오버플로우 위험 없이 안전하게 동작한다는 장점이 있습니다.