Kth Smallest Element in a BST

The problem description can be found in LeetCode.

Approach

Since we are looking for the k-th smallest element in a Binary Search Tree, the smallest ones are on the left side.

In particular, if we follow the In-order traversal, we would get an ordered list of values.

In this traversal, the order is:

1. Left.
2. Node.
3. Right.

We can use the 2. step to check the value of the node and keep track of the answer in a global variable, which we can update as soon as we found the k-th smallest element.

Here, we can either start from k and go until 0 (as implemented) or start from 0 and go until k.

Solution

from typing import Optional


class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def dfs(node: Optional[TreeNode]) -> None:
            nonlocal ans
            nonlocal i

            if not node:
                return

            dfs(node.left)

            i -= 1
            if i == 0:
                ans = node.val
            else:
                dfs(node.right)

        ans = -1
        i = k

        dfs(root)

        assert ans != -1, "expected one solution"
        return ans

Complexity

• Time: $O(n)$
• Space: $O(h)$

Where $n$ is the number of nodes in the tree and $h$ is its height. If the tree is unbalanced, it is like having a Linked List, hence $h=n$.