Problem overview
The method KthSmallest receives the root of a binary search tree of TreeNode objects and an integer k. It returns the k-th smallest value in the tree, counting from 1, so k = 1 asks for the smallest value. All values are different, and k is between 1 and the number of nodes.
For the tree with root 3, left child 1 (which has a right child 2) and right child 4, the values in increasing order are 1, 2, 3, 4. With k = 1 the answer is 1. For the tree with root 5, children 3 and 6, then 2 and 4 under 3, and 1 under 2, the sorted values are 1, 2, 3, 4, 5, 6, so k = 3 gives 3 and k = 6 gives 6.
Approach: in-order traversal that stops at the k-th node
In a binary search tree, everything in a node's left subtree is smaller than the node, and everything in its right subtree is larger. An in-order traversal (the left subtree, then the node, then the right subtree) therefore visits the values in increasing order. The k-th node it visits holds the answer.
The simplest way to use that is to collect every value into a list and take the element at index k - 1. It works, but it visits the whole tree even when k is small. The follow-up asks to stop as soon as the answer is known, in O(h + k) time, where h is the height.
A recursive traversal can't just return in the middle of its calls, so the lesson writes the traversal with an explicit stack, which can stop at any moment. The stack holds the nodes passed on the way down that still wait to be visited:
- Go as far left as possible, pushing every node on the way. The last node pushed is the smallest one not yet visited.
- Pop a node. This is the next value in increasing order. Count it.
- If the count is
k, return its value. - Otherwise move to its right child and repeat. The right subtree holds values larger than the popped node but smaller than the nodes still on the stack.
On the tree with root 5 and k = 3: the walk pushes 5, 3, 2 and 1. It pops 1 (count 1) and finds no right child, then pops 2 (count 2), then pops 3 (count 3) and returns 3. The nodes 4, 5 and 6 are never visited.
Step by step
- Create the stack.
stackis a new, emptyStack<TreeNode>.Pushputs a node on top andPoptakes the top one off. - Create
node, set toroot. It is the part of the tree still to walk through. - Create
countat 0: how many nodes have been visited. - Outer loop. Run while
node != nullor the stack is not empty (stack.Count > 0). The traversal is over only when there is nothing to walk into and nothing waiting. - Inner loop. While
node != null, head to the leftmost node, since smaller values are on the left. - Push
nodeinside the inner loop. It waits there until its left subtree has been visited. - Move left with
node = node.left. - Pop the top node into
nodeonce the inner loop ends. It is the smallest value not yet visited. - Visit it with
count++. - Check
count == k. If so, returnnode.val. - Move right with
node = node.right. The next pass walks through that subtree. - Return -1 after the loop. Since
kis at most the number of nodes, the loop returns first. C# requires a return on every path.
The solution in C#
class Solution {
public int KthSmallest(TreeNode root, int k) {
Stack<TreeNode> stack = new Stack<TreeNode>();
TreeNode node = root;
int count = 0;
while (node != null || stack.Count > 0) {
while (node != null) {
stack.Push(node);
node = node.left;
}
node = stack.Pop();
count++;
if (count == k) {
return node.val;
}
node = node.right;
}
return -1;
}
}
The C# details worth knowing:
Stack<T>lives inSystem.Collections.Generic. A standalone file needsusing System.Collections.Generic;. Its size is the propertyCount, notLength.Popthrows on an empty stack. It raises anInvalidOperationException. The loop condition guarantees that, after the inner loop, the stack holds a node: either it was non-empty to begin with, ornodewas non-null and has just been pushed. SoPopis safe here.nodedoes double duty. It is the cursor while walking left, and it holds the popped node afterwards. Assigning to it never changes the tree, only where the variable points.count++as a statement adds 1 tocount. Its old value isn't used here, so++countwould do the same.- Why the final
return -1;is there. The compiler can't prove that the loop always returns, and a method returningintmust return on every path.
Complexity
Let h be the height of the tree and n the number of nodes.
Time: O(h + k). The first descent pushes at most h nodes. After that, every pop counts one more node in sorted order, and the method stops at the k-th pop. Every node is pushed once before it is popped, so the number of pushes is at most the number of pops plus the nodes still on the stack: at most k + h. In the worst case, k = n, this becomes O(n).
Space: O(h). The stack holds the path from the root to the current node: about log n nodes for a balanced tree and up to n for a chain. Because the traversal is iterative, a deep tree uses heap memory for the stack and doesn't risk exhausting the call stack, as deep recursion could.
Edge cases and common mistakes
- One node. With a single node and
k = 1, the walk pushes it, pops it, counts 1 and returns its value. - The smallest and the largest. In the second example tree,
k = 1returns 1 andk = 6returns 6. - A chain going left. If every node has only a left child, the first descent pushes them all.
k = 1returns after that single descent, andk = nreturns the root, which is the largest value. - A right child in the middle. In the first example tree, the node 1 has only a right child. After popping 1, the walk moves right to 2, which is the correct next value. Skipping the move to
node.rightwould lose it. - Counting from 0. The statement counts from 1. Starting
countat 1, or comparing withk - 1without adjusting, returns the wrong node. - Forgetting the second part of the loop condition. With only
node != null, the loop would end as soon asnodebecomesnullafter a leaf, while nodes are still waiting on the stack.
Testing
- Tree
[3, 1, 4, null, 2]withk = 1gives1, withk = 2gives2and withk = 4gives4. - Tree
[5, 3, 6, 2, 4, null, null, 1]withk = 3gives3. - The same tree with
k = 6gives6. - A one-node tree with
k = 1gives that node's value.
