Problem overview
The method kthSmallest receives the root of a binary search tree 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 distinct, and k is between 1 and the number of nodes.
For the tree with root 5, children 3 and 6, then 2 and 4 under 3, and 1 under 2, the values in increasing order are 1, 2, 3, 4, 5, 6. With k = 3 the answer is 3, and with k = 6 it is 6. For the tree [3,1,4,null,2] and k = 1, the answer is 1.
The follow-up asks you to stop as soon as the answer is found, in O(h + k) time, where h is the height of the tree.
Approach: in-order traversal that stops at the k-th node
In a BST every value in the left subtree is smaller than the node and every value in the right subtree is larger. An in-order traversal visits the left subtree, then the node, then the right subtree, so it visits the values in increasing order. The k-th node it visits holds the answer.
Collecting all the values into a list and indexing it would work, but it always visits every node. Counting visits and stopping at the k-th one avoids that. Doing that with plain recursion would need a way to stop all the pending calls, so this lesson uses an explicit stack instead. The loop can then return the moment the count reaches k.
The stack holds the nodes we passed on the way down and haven't visited yet. The idea runs in three moves:
- From the current node, go left as far as possible, pushing every node on the way. The last one is the smallest value in this part of the tree.
- Pop the top node. Everything smaller than it has been visited, so visit it: count it, and return its value if the count is
k. - Otherwise move to its right child and repeat. The right subtree holds values larger than this node but smaller than the nodes still waiting on the stack.
Step by step
- Create a
Stack<TreeNode>namedstack. It holds the nodes whose own value and right subtree are still to come. - Start
nodeatroot. It is the root of the part of the tree still to be walked through. - Declare
count = 0. It counts the nodes visited so far. - Loop while
node != null || !stack.isEmpty(). The walk is finished only when no subtree is left to walk and no node is waiting on the stack. - Inner loop, while
node != null: head to the leftmost node, because smaller values are to the left. - Push
nodeinside the inner loop. Its left subtree has smaller values, so it waits. - Move to
node.leftand let the inner loop continue. - After the inner loop, pop the top node into
node. It is the smallest value not yet visited. - Count the visit with
count++. - If
count == k, returnnode.val. The rest of the tree is never touched. - Otherwise go to
node.right. The next pass of the outer loop walks that subtree. - Return -1 after the outer loop. Since
kis at most the number of nodes, the loop returns first. The line exists because Java wants a return on every path.
The solution in Java
class Solution {
public int kthSmallest(TreeNode root, int k) {
Stack<TreeNode> stack = new Stack<>();
TreeNode node = root;
int count = 0;
while (node != null || !stack.isEmpty()) {
while (node != null) {
stack.push(node);
node = node.left;
}
node = stack.pop();
count++;
if (count == k) {
return node.val;
}
node = node.right;
}
return -1;
}
}
Java details worth a closer look:
Stack<TreeNode>and the diamond.new Stack<>()lets the compiler infer the type argument from the left side. With the type given,pop()returns aTreeNode. A rawStackwould returnObjectand the assignment tonodewouldn't compile.- The class lives in
java.util. The code usesStackwithout an import line, so the lesson already hasjava.utilavailable. In a file of your own, addimport java.util.*;. pop()on an empty stack throwsEmptyStackException. The outer condition makes sure the stack has something whenevernodeisnull, so the pop is safe.- Reusing
node. The same variable is the walker going down, the popped node and the move to the right child. Follow it carefully when you trace the code. - Newer code often uses
ArrayDeque<TreeNode>withpushandpopinstead ofStack. It has the same calls here and avoids the old synchronisedVectorunderneath.
Complexity
Time: O(h + k). Reaching the smallest value takes h pushes down the left spine. After that, each visited node costs a pop, a count and possibly a few pushes, and every node is pushed and popped at most once. Stopping after k visits means that about k nodes are processed beyond the first descent. Even without the early return, the loop would end after at most n visits, so the worst case is O(n).
Space: O(h). The stack holds one path from the root down, at most h nodes: O(log n) in a balanced tree and O(n) in a tree that is a single left-leaning chain.
Edge cases and common mistakes
- A single node. With
k = 1, the node is pushed, popped, counted and returned. kequals the number of nodes. The loop visits every node and returns the largest value on the last visit, which is whyk = 6returns 6 in the second example.- Counting before popping. Incrementing
countwhile pushing counts nodes in the wrong order. Only the pop is a visit. - Forgetting the right subtree. Without
node = node.right, the walk visits only the left spine. In the first example,[3,1,4,null,2], the value 2 sits to the right of 1, andk = 2would then return the wrong value. - Using only
node != nullas the loop condition. After a leaf is popped,nodebecomesnullwhile nodes still wait on the stack, so the loop would end early. - Counting from 0.
k = 1means the smallest value, so the check iscount == kafter the increment.
Testing
- Tree
[3,1,4,null,2]withk = 1gives1. - Tree
[5,3,6,2,4,null,null,1]withk = 3gives3. - The same tree with
k = 6gives6. - The same tree with
k = 1gives1.
