Problem overview
The function kth_smallest receives the root of a binary search tree and an integer k, and returns the k-th smallest value in the tree. Counting starts at 1, so k = 1 asks for the minimum. All values are different, and k is between 1 and the number of nodes.
For the tree [3, 1, 4, None, 2] and k = 1, the values in increasing order are 1, 2, 3, 4, so the answer is 1. For the tree [5, 3, 6, 2, 4, None, None, 1] the sorted values are 1, 2, 3, 4, 5, 6: k = 3 gives 3 and k = 6 gives 6.
Approach: walk the tree in order and stop at the k-th node
In a binary search tree, an in-order traversal, which visits the left subtree, then the node, then the right subtree, produces the values in increasing order. So the k-th node visited holds the answer.
The simple way is to collect every value into a vector during a recursive traversal and then read index k - 1. That visits every node even when k is 1. The follow-up asks for O(h + k) time, where h is the tree's height: go down to the smallest value, then visit only k nodes. To stop in the middle of a traversal cleanly, this lesson writes the traversal with an explicit stack instead of recursion, and returns the moment the count reaches k.
The stack holds the nodes whose left subtrees are still being visited. Here is k = 3 on the second tree. The walk goes left from 5 to 3 to 2 to 1, pushing each, so the stack is 5, 3, 2, 1. It pops 1 (the first visit) and 1 has no right child. It pops 2 (the second visit), which has no right child either. It pops 3, the third visit, and returns 3. The nodes 4, 5 and 6 were never looked at, apart from 5 sitting on the stack.
Step by step
- Create the stack. A
Vec<&TreeNode>holds borrowed nodes.pushandpopwork on the end of the vector, which is all a stack needs. - Create the cursor.
nodestarts asroot.as_deref(), anOption<&TreeNode>pointing at the part of the tree that still has to be walked. - Count visits.
countstarts at 0. - Outer loop. It runs while the cursor holds a node or the stack isn't empty. The traversal is over only when both are exhausted.
- Inner loop: go left.
while let Some(current) = nodepushescurrenton the stack and movesnodetocurrent.left.as_deref(). Each node waits on the stack until everything smaller, its left subtree, has been visited. - Pop the next node. When the inner loop ends, the cursor is
None, and the top of the stack is the smallest unvisited node.pop()returns anOption, andunwrap()is safe because the stack can't be empty at this point. - Visit it. Add 1 to
count. Ifcount == k, returncurrent.val. - Continue with the right subtree. Set
nodetocurrent.right.as_deref(). Its values are bigger than the popped node but smaller than the nodes still on the stack, so they come next. - Add a tail value.
-1without a semicolon after the loops gives the function a value on every path. It is never reached, sincekis at most the number of nodes.
The solution in Rust
impl Solution {
fn kth_smallest(root: Option<Box<TreeNode>>, k: i32) -> i32 {
let mut stack: Vec<&TreeNode> = Vec::new();
let mut node = root.as_deref();
let mut count = 0;
while node.is_some() || !stack.is_empty() {
while let Some(current) = node {
stack.push(current);
node = current.left.as_deref();
}
let current = stack.pop().unwrap();
count += 1;
if count == k {
return current.val;
}
node = current.right.as_deref();
}
-1
}
}
Rust details worth noticing:
- References on the stack. The stack holds
&TreeNode, not boxes, so pushing a node doesn't move or copy it. The references point into the tree owned byroot, which stays alive until the function returns. - Two variables called
current. The innerwhile letbinds one that exists only inside that loop. Thelet current = stack.pop().unwrap();after it is a new variable, so there is no clash. - Type of
count. It is compared withk, ani32, so Rust infersi32for it fromcount == k. unwrap()onpop(). It would panic on an empty stack. Here the inner loop has just pushed a node, or the stack already held one, so it can't happen.- Returning from inside loops.
return current.val;exits both loops and the function in one go, which is why the traversal can stop early.
Complexity
Time: O(h + k). The first descent pushes at most h nodes. After that, every pop is one visit, and the code stops at visit k. Each node is pushed and popped at most once, and the pushes are bounded by the pops plus the height. In the worst case, k = n, that is O(n).
Space: O(h). The stack holds a path from the root down, so it never has more than h nodes: about log n for a balanced tree and up to n for a tree that is one long chain.
Edge cases and common mistakes
k = 1. The answer is the leftmost node. The first descent finds it and the first pop returns it.kequal to the node count. The traversal visits everything and returns the largest value.- A single node.
k = 1returns the root's value. - A chain that only goes right. The inner loop pushes one node, the pop visits it, and the cursor moves right. The stack stays tiny.
- Counting before the pop. Counting nodes as they are pushed counts them in pre-order, not in sorted order. Count when you pop.
- Forgetting the right subtree. Without
node = current.right.as_deref(), the traversal skips every node that is only reachable through a right child, and the count comes out wrong. - Testing only
node.is_some()in the outer loop. The loop would stop while nodes still wait on the stack.
Testing
[3, 1, 4, None, 2]withk = 1gives1.[5, 3, 6, 2, 4, None, None, 1]withk = 3gives3.[5, 3, 6, 2, 4, None, None, 1]withk = 6gives6.[5, 3, 6, 2, 4, None, None, 1]withk = 1gives1.[5, 3, 6, 2, 4, None, None, 1]withk = 2gives2.
