Problem overview
The function kthSmallest receives the root of a binary search tree and an integer k counted from 1, and returns the k-th smallest value in the tree. The values are distinct and 1 <= k <= n, where n is the number of nodes.
For the tree with root 5, children 3 and 6, where 3 has children 2 and 4 and 2 has a left child 1, the sorted values are 1, 2, 3, 4, 5, 6. So k = 3 gives 3 and k = 6 gives 6. For the tree [3,1,4,NULL,2] with k = 1 the answer is 1.
Approach: in-order traversal that stops early
An in-order traversal visits the left subtree, then the node, then the right subtree. In a BST that visits the values in increasing order, so the k-th node visited holds the answer.
A recursive traversal is the usual way to write it, but it makes stopping at the k-th node awkward. The lesson uses an explicit stack instead, which turns the traversal into a loop that can simply return. The stack holds the nodes on the way down whose left subtree is still being visited. This also gives the follow-up's bound: reaching the first node costs about h pushes, where h is the height, and each further visit costs a pop and a few pushes, so the work is O(h + k) rather than O(n).
On the second tree with k = 3: the walk pushes 5, 3, 2, 1 going left. It pops 1 (count 1), has no right child, pops 2 (count 2), pops 3 (count 3) and returns 3. Nodes 4, 5 and 6 are never visited.
Step by step
- Create an empty stack of node pointers. It holds the nodes passed on the way down that wait to be visited.
- Set
nodeto the root. It is the root of the part of the tree that is still to be walked. - Count visits with
count, starting at 0. - Loop while there is work left:
nodeis not null, or the stack is not empty. The traversal is finished only when neither holds. - Walk to the leftmost node with an inner loop that runs while
nodeis not null, because the smallest value of any subtree is at its leftmost node. - Push each node on the way down. Its left side holds smaller values, so it has to wait.
- Move to the left child.
- Take the next node to visit. After the inner loop,
nodeis null. Read the top of the stack withtop(), then remove it withpop(). These are two calls, becausepop()returns nothing. - Count the visit with
count++. - Check for the answer. If
count == k, return this node's value. - Otherwise go right. The next larger values are in this node's right subtree, so set
node = node->rightand let the next pass of the outer loop handle it. - Return -1 after the loop. Since
kis at mostn, the function always returns earlier. The line just gives the function a return value on every path.
The solution in C++
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
stack<TreeNode*> stack;
TreeNode* node = root;
int count = 0;
while (node != NULL || !stack.empty()) {
while (node != NULL) {
stack.push(node);
node = node->left;
}
node = stack.top();
stack.pop();
count++;
if (count == k) {
return node->val;
}
node = node->right;
}
return -1;
}
};
The C++ details worth a closer look:
std::stack.pushadds on top,topreads the top element,popremoves it without returning it, andemptytests for no elements. By default it is built on adeque.top()beforepop(). Once an element is popped, a reference to it is invalid. The code copies the pointer intonodefirst, then pops.- Naming the stack
stack. The variable has the same name as the type. That is legal here, because the declarationstack<TreeNode*> stack;is read with the type first, but inside the function the name now refers to the variable. Another name, such asst, avoids confusion. - Calling
top()on an empty stack is undefined behaviour. Here it can't happen: the outer loop only runs whilenodeis non-null or the stack is not empty, and a non-nullnodeis pushed by the inner loop first. - Pointers on the stack. The stack stores copies of the pointers, not the nodes, so pushing is cheap and the tree is not modified.
Complexity
Time: O(h + k). The first descent pushes at most h nodes. After that, every visit pops one node, and the pushes of the later descents add up to at most the nodes that are visited or still waiting, so the whole run does O(h + k) pushes and pops. In the worst case, k = n, this is O(n).
Space: O(h). The stack holds at most one path from the root down, which is h nodes: about log n for a balanced tree and n for a chain.
Edge cases and common mistakes
k = 1. The answer is the leftmost node. The first pop returns it.k = n. The answer is the largest value, and the whole tree is visited.- A single node.
kcan only be 1, and the root is returned. - No left children. For a chain that goes right, the stack never holds more than one node, and the values come out in order anyway.
- Forgetting
node = node->right. Without it, the same node would be pushed again and the loop would never end, or it would skip right subtrees. - Counting before the walk to the left. The count must go up when a node is visited (popped), not when it is pushed. Pushing happens before the smaller values have been seen.
- Treating
kas 0-based. The first smallest value isk = 1.
Testing
[3,1,4,NULL,2]withk = 1gives 1.[5,3,6,2,4,NULL,NULL,1]withk = 3gives 3.[5,3,6,2,4,NULL,NULL,1]withk = 6gives 6.
