Problem overview
The method kthSmallest receives the root of a binary search tree and a number k, counted from 1, and returns the k-th smallest value in the tree. All values are different, and k is between 1 and the number of nodes.
For the tree [3, 1, 4, null, 2] the values in increasing order are 1, 2, 3, 4, so k = 1 gives 1. For [5, 3, 6, 2, 4, null, null, 1] the order is 1, 2, 3, 4, 5, 6, so k = 3 gives 3 and k = 6 gives 6.
Approach: in-order traversal, stopped early
In a binary search tree, an in-order traversal (the left subtree, then the node, then the right subtree) visits the values in increasing order. So the k-th node the traversal visits holds the answer, and the traversal can stop right there instead of visiting the rest.
The recursive version is short, but stopping in the middle of a recursion needs extra bookkeeping. Here the traversal uses an explicit stack instead:
- From the current node, go left as far as possible, pushing every node on the way. These nodes wait: they are bigger than everything in their left subtree.
- Pop the top node. It is the smallest value not visited yet. Count it.
- If the count equals
k, return its value. Otherwise continue with its right subtree, whose values come next in order.
On the second example with k = 3: starting at 5, the walk pushes 5, 3, 2 and 1, then reaches null. It pops 1 (count 1) and its right side is empty, so the next pop is 2 (count 2). It has no right child either, so the next pop is 3 (count 3), which equals k. The answer is 3, and 4 and 6 were never visited.
Step by step
- Create
stackas aconst, typedTreeNode[]. It holds nodes that were passed on the way down and not visited yet. - Create
nodewithletand set it toroot. It is the root of the part of the tree still to walk. - Create
countwithlet, starting at 0, the number of visited nodes. - Loop while
node || stack.length > 0: there is a subtree to walk or a node waiting on the stack. Only when both are exhausted is the traversal over. - Go left with an inner loop that runs while
nodeis notnull. - Push
nodeinside it. It waits until its left subtree has been visited. - Move to
node.left. - Pop the top node into
node, writing!afterpop(). This is the smallest unvisited value. - Count the visit with
count++. - Return
node.valifcount === k. - Otherwise move to
node.right. Its values are larger than this node but smaller than the nodes still on the stack. - Return -1 after the loop. Since
kis at most the number of nodes, the loop returns first, and the line only keeps the return type honest.
The solution in TypeScript
class Solution {
kthSmallest(root: TreeNode | null, k: number): number {
const stack: TreeNode[] = [];
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;
}
}
The TypeScript details worth a look:
- The
!afterpop().Array.prototype.popreturnsundefinedfor an empty array, so with strict null checking its type isTreeNode | undefined. The non-null assertion!tells the compiler that the stack isn't empty here. It is erased in the JavaScript output and checks nothing at run time. It is true here because the inner loop has just pushed a node, or the outer condition guaranteed one on the stack. - Truthiness as a null check.
while (node)andnode || stack.length > 0rely on objects being truthy andnullfalsy. That is safe for nodes. It would not be for numbers, where0is falsy. - One variable, two jobs.
nodeis first the subtree to walk, then the node just popped, then its right child. Its inferred type,TreeNode | null, covers all three. - A
conststack still changes.constonly stops the variable from being reassigned.pushandpopmodify the array in place.
Complexity
Let h be the height of the tree.
Time: O(h + k). The first descent pushes at most h nodes. After that, the loop pops one node per visit, and each pop is paired with the pushes of its left spine, so the total work to reach the k-th node is bounded by h + k. Even a full traversal pushes and pops each node at most once, so it is O(n).
Space: O(h). The stack holds the nodes along one path from the root. For a tree shaped like a list that is up to n nodes.
Edge cases and common mistakes
- A single node.
[7]withk = 1pushes 7, pops it, counts 1 and returns 7. kequal to the number of nodes. For the six-node tree,k = 6visits everything and returns 6, the largest value.- A node with only a right child. In
[3, 1, 4, null, 2]the node 1 has only the child 2. After popping 1, the walk goes tonode.rightand finds 2 before returning to 3. Withk = 4the answer is 4. - Counting from 0.
k = 1means the smallest value, so the check iscount === kafter incrementing. - Using a pre-order walk. Visiting the node before its left subtree gives the values in tree order, not in sorted order.
- Forgetting the right subtree. Without
node = node.right, the loop would keep popping the same stack and miss every value that is in a right subtree.
Testing
[3, 1, 4, null, 2]withk = 1gives1.[3, 1, 4, null, 2]withk = 4gives4.[5, 3, 6, 2, 4, null, null, 1]withk = 3gives3.[5, 3, 6, 2, 4, null, null, 1]withk = 6gives6.[7]withk = 1gives7.
