Kth Smallest Element in a BST

How it works

Inorder traversal of a BST gives values in sorted order, so the k-th visited node is the k-th smallest. But recursion and stacks use extra space.

Morris Traversal allows us to perform inorder traversal using O(1) extra space, by temporarily creating a "thread" (a right pointer) from a node's predecessor back to the node.

For each node:

• If it has no left child - visit it directly.
• If it has a left child - find its inorder predecessor.
  • If the predecessor's right pointer is empty - create a temporary link to the current node and move left.
  • If the predecessor's right pointer already points to the current node - remove the link, visit the node, and move right.

We decrement k each time we "visit" a node. The node where k becomes 0 is the k-th smallest.

This works because we simulate the inorder order without extra memory.

Steps

1. Set curr = root.
2. While curr is not null:
   • Case 1: No left child
     • Visit curr (decrement k).
     • If k == 0, return curr.val.
     • Move to curr.right.
   • Case 2: Has a left child
     • Find the inorder predecessor (pred = rightmost node in left subtree).
     • If pred.right is null:
       • Create a temporary thread: pred.right = curr.
       • Move curr to its left child.
     • Else (thread already exists):
       • Remove the thread: pred.right = null.
       • Visit curr (decrement k).
       • If k == 0, return curr.val.
       • Move to curr.right.
3. If traversal ends without finding k nodes, return -1.

Complexity

• Time complexity: O(n)
• Space complexity: O(1)

Things to watch for

Using Pre-Order or Post-Order Instead of In-Order

A BST's defining property is that in-order traversal yields sorted values. Using pre-order or post-order traversal does not produce sorted output, so you cannot simply pick the k-th visited node. Always use in-order traversal (left, node, right) to get elements in ascending order.

Off-by-One Errors with k

The k-th smallest element is at index k-1 in a zero-indexed array or corresponds to the k-th decrement of a counter starting at k. Mixing up one-based and zero-based indexing leads to returning the wrong element. Be consistent: if using a counter, return when it reaches zero; if using an array, access index k-1.

Not Stopping Early

When counting nodes during traversal, there is no need to visit the entire tree once the k-th element is found. Failing to return early means wasted computation, especially in large trees where k is small. Use a flag or check the counter after each visit to terminate as soon as possible.

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:
        curr = root

        while curr:
            if not curr.left:
                k -= 1
                if k == 0:
                    return curr.val
                curr = curr.right
            else:
                pred = curr.left
                while pred.right and pred.right != curr:
                    pred = pred.right

                if not pred.right:
                    pred.right = curr
                    curr = curr.left
                else:
                    pred.right = None
                    k -= 1
                    if k == 0:
                        return curr.val
                    curr = curr.right

        return -1

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        TreeNode* curr = root;

        while (curr) {
            if (!curr->left) {
                k--;
                if (k == 0) return curr->val;
                curr = curr->right;
            } else {
                TreeNode* pred = curr->left;
                while (pred->right && pred->right != curr)
                    pred = pred->right;

                if (!pred->right) {
                    pred->right = curr;
                    curr = curr->left;
                } else {
                    pred->right = nullptr;
                    k--;
                    if (k == 0) return curr->val;
                    curr = curr->right;
                }
            }
        }
        return -1;
    }
};