BST가 주어질 때 K번째 작은 원소를 구하는 문제

Solution

BST를 탐색하면서 K번째 작은 원소가 나오면 리턴하면 된다.

BST 탐색 방법 중 inorder 탐색을 하면(왼쪽 서브트리, 자신, 오른쪽 서브트리) 크기 순으로 탐색할 수 있고 이를 배열에 넣은 다음에 배열에서 k-1 인덱스의 원소를 뽑아내면 된다.

중간에 k번 탐색을 했다면 바로 리턴하는 식으로 가지치기는 할 수 있으나 어차피 O(n) 이므로 큰 의미는 없다.

Source Code

from typing import Optional


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        nodes = []

        def inorder(node: TreeNode):
            if not node:
                return

            inorder(node.left)
            nodes.append(node.val)
            inorder(node.right)

        inorder(root)
        return nodes[k - 1]

230. Kth Smallest Element in a BST

Problem link

Problem Summary

Solution

Source Code