Problem overview
The method kthSmallest receives the root of a binary search tree and an integer k, counted from 1. It returns the k-th smallest value in the tree, so k = 1 asks for the minimum. All values are different, 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, the largest value. For the tree with root 3, left child 1 (which has 2 as its right child) and right child 4, the sorted values are 1, 2, 3, 4, and k = 1 gives 1.
The follow-up asks you to stop as soon as the answer is known, in O(h + k) time, where h is the height of the tree.
Approach: in-order traversal, stopped early
In a BST, everything in a node's left subtree is smaller than the node and everything in its 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.
A recursive in-order traversal would collect every value, or need extra work to stop in the middle. This lesson uses an explicit stack instead, which makes stopping natural: the loop simply returns when the count reaches k, and the rest of the tree is never touched.
The stack holds the nodes on the way down that are waiting for their turn. From any node, the smallest value in its part of the tree is at its leftmost node. So the walk goes left as far as possible, pushing each node it passes. Then it pops the top node, which is the smallest one not yet visited, counts it, and continues with that node's right subtree, whose values come next in order.
Take the tree with root 3, left child 1 (right child 2) and right child 4, and k = 3. Going left from 3 pushes 3, then 1. Pop 1: count is 1. Its right child is 2, so push 2 and pop it: count is 2. It has no right child, so the next pop gives 3: count is 3, which equals k, and the method returns 3.
Step by step
- Create
stackas an empty array, withconst, since the variable is never reassigned.pushandpopboth work at the end of an array, so an array serves as a stack. - Declare
nodewithletand set it toroot. It is the root of the part of the tree still to be walked through. When it isnull, the next node to visit is the one on top of the stack. - Declare
countwithlet, starting at 0. It counts the nodes visited so far. - Write the outer
whileloop with the conditionnode || stack.length > 0. The traversal is finished only when nothing remains to walk through and no node is waiting on the stack. - Write an inner
while (node)loop to head for the leftmost node, because each step to the left leads to smaller values. - Push
nodeonto the stack. Its left subtree holds smaller values, so it must wait until they have been visited. - Move
nodetonode.left. This ends the inner loop when the walk falls off the left edge. - After the inner loop, pop the top of the stack into
node. This is the smallest value not yet visited: everything smaller, including its whole left subtree, has been visited already. - Visit it with
count++. - If
count === k, returnnode.val. This is the early exit that givesO(h + k). - Otherwise set
nodetonode.right. The next values in order are in the right subtree. They are larger than this node and smaller than whatever is still on the stack, so the next pass walks through them first. - After the outer loop, return -1. Since
kis at most the number of nodes, the method returns earlier. The line makes sure a number is returned on every path.
The solution in JavaScript
class Solution {
kthSmallest(root, k) {
const stack = [];
let node = root;
let count = 0;
while (node || stack.length > 0) {
while (node) {
stack.push(node);
node = node.left;
}
node = stack.pop();
count++;
if (count === k) {
return node.val;
}
node = node.right;
}
return -1;
}
}Some details about JavaScript:
constdoesn't make the array immutable.const stack = []only stops the namestackfrom being rebound. Pushing and popping still change the array.pop()returnsundefinedon an empty array. Here it never does: the outer condition is true on a pass only ifnodeis a node or the stack is not empty, and the inner loop pushes at least one node whenevernodeis a node. So the stack is never empty at the pop.node || stack.length > 0relies on truthiness. A node is an object, so it is truthy, andnullis falsy. A node whosevalis 0 is still truthy.===comparescountandkwithout conversion. Both are numbers, so it behaves like the comparison you expect.
Complexity
Going down to the first leftmost node pushes at most h nodes. After that, each node visited is popped once, and each push is paid for by a later pop, so reaching the k-th node takes O(h + k) steps. Even without the early return, every node is pushed and popped at most once, so the whole traversal is O(n).
The stack holds one path from the root, so the extra space is O(h): about log n for a balanced tree and n for a chain.
Edge cases and common mistakes
- Counting from 0.
kstarts at 1. Checking the count before incrementing it would return the wrong node. - Forgetting the right subtree. Without
node = node.rightafter the visit, the loop would pop the next stack node and skip every right child. On the tree with root 3, left child 1 and 1's right child 2, the second smallest value, 2, would be missed. - Stopping the outer loop when
nodeisnull. After a pop,nodemay benullwhile the stack still holds nodes. That is why the condition also checks the stack. - A single-node tree. With
k = 1, the walk pushes the root, pops it, counts 1 and returns its value. - A chain tilted to one side. The stack can grow to
nnodes when the tree only has left children. Right-only chains keep the stack at one node.
Testing
On the tree with root 5, children 3 and 6, 2 and 4 under 3, and 1 under 2, kthSmallest(root, k) returns 1, 2, 3, 4, 5 and 6 for k from 1 to 6.
