Kth Smallest Element in a BST

Problem

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note: You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Follow up: What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Solution

Recursive Solution

public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        traversal(root, k);
        return result;
    }

    int level = 1, result = 0;

    private void traversal(TreeNode node, int k) {
        if (node == null) return;
        traversal(node.left, k);
        if (level++ == k) result = node.val;
        else traversal(node.right, k);
    }
}

Iterative Solution using Stack

public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> s = new Stack<>();
        TreeNode node = root;
        int level = 1;
        while (!s.isEmpty() || node != null) {
            if (node != null) {
                s.push(node);
                node = node.left;
            }
            else {
                node = s.pop();
                if (level++ == k) return node.val;
                node = node.right;
            }
        }
        return 0;
    }
}

Analysis

The BST has values in order node.left -> node -> node.right
Hence it's obvious that we need to have inorder traversal
Because the most-left node is smallest, we won't update level until we reach the end of left subtree
Then we check if level reaches k, if so, return current node.val, otherwise, we turn to node.right
We increment level before we go to node.right