Problem overview
The method kthSmallest receives the root of a binary search tree and an integer k, and returns the k-th smallest value in the tree. Counting starts at 1, so k = 1 asks for the minimum. The values are all different, and k is between 1 and the number of nodes.
For the tree [3, 1, 4, null, 2], with root 3, children 1 and 4, and 2 as the right child of 1, the values in increasing order are 1, 2, 3, 4. With k = 1 the answer is 1. For the tree [5, 3, 6, 2, 4, null, null, 1] the values in order are 1, 2, 3, 4, 5, 6, so k = 3 gives 3 and k = 6 gives 6.
Approach: visit the nodes in increasing order and stop at the k-th
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 from smallest to largest. The k-th node it visits holds the answer.
A recursive traversal is short, but it is awkward to stop early from deep inside recursion. The lesson uses an explicit stack instead, which makes the traversal a plain loop that can return the moment the count reaches k. This matches the follow-up: O(h + k) time, where h is the height of the tree. The traversal descends to the leftmost node, which takes up to h steps, and then visits k nodes in total.
The stack holds the nodes on the way down that still wait for their turn. The loop repeats three actions:
- From the current node, go left as far as possible, pushing every node on the way. The leftmost node holds the smallest value of this part of the tree.
- Pop the top node. Everything smaller than it has been visited. Count it. If the count is
k, it is the answer. - Otherwise move to its right child. That subtree holds the next larger values, bigger than the popped node but smaller than the nodes still waiting on the stack. The loop then repeats.
Trace the first example with k = 1. Push 3, go left, push 1, and go left again to null. Pop 1: the count becomes 1, which equals k, so the answer is 1. The nodes 4 and 2 are never looked at.
Step by step
- Create the stack in a
val: an emptyscala.collection.mutable.Stack[TreeNode]().pushputs a node on top, andpoptakes the top one off. - Declare
nodewithvar, starting atroot. It is the root of the part of the tree that remains to be walked. - Declare
countwithvar, starting at 0. It is the number of nodes visited so far. - Write the outer loop that runs while
node != null || stack.nonEmpty. There is work left if there is a subtree to walk or a node waiting on the stack. - Write an inner loop that runs while
node != null, to go left. - Push
nodeinside it. It waits there until its left subtree has been visited. - Move to
node.left. - Pop the top after the inner loop with
node = stack.pop(). It is the smallest value not yet visited. - Count the visit with
count += 1. - Return when
count == k:return node.value. The rest of the tree is never touched. - Otherwise go right with
node = node.right, so the next pass walks that subtree. - Return -1 after the loop. Since
kis at most the number of nodes, this line is never reached. The method still has to produce anInton every path.
The solution in Scala
class Solution {
def kthSmallest(root: TreeNode, k: Int): Int = {
val stack = scala.collection.mutable.Stack[TreeNode]()
var node = root
var count = 0
while (node != null || stack.nonEmpty) {
while (node != null) {
stack.push(node)
node = node.left
}
node = stack.pop()
count += 1
if (count == k) {
return node.value
}
node = node.right
}
return -1
}
}
The Scala details worth knowing:
mutable.Stack. The collection is inscala.collection.mutable.pushadds on top,pop()removes and returns the top element, andnonEmptytells whether anything is left. Callingpop()on an empty stack throwsNoSuchElementException. In this loop the outer condition guarantees that the stack is non-empty whenpop()runs, because the inner loop either pushed something or the stack was non-empty to begin with.- The type argument.
Stack[TreeNode]()creates an empty stack ofTreeNodes. Without the type, the compiler would have nothing to infer it from. var nodeandval stack. The stack object is changed in place and the variable never points elsewhere, sovalis right.nodeis reassigned on every step, so it is avar.||in the loop condition. It means "keep going while there is a subtree to walk or a node waiting". With&&the loop would stop too early, because at the start the stack is empty.returninside nested loops. It leaves the whole method at once, not only the loop it is in. That is what the early stop needs.- Counting from 1.
countis incremented before the comparison, so the first visited node, the smallest, makes it 1 and answersk = 1.
Complexity
Time: O(h + k). The first descent costs up to h pushes, and after that each visited node costs a pop plus the pushes of its right subtree's left spine. Every node is pushed once and popped at most once, and the loop stops after k pops. Without the early return, the whole traversal would be O(n).
Space: O(h). The stack holds at most one root-to-node path, which is the height of the tree. For a balanced tree that is about log n. For a tree that is a single left-leaning chain it is n.
Edge cases and common mistakes
- A single node. The only valid
kis 1. The node is pushed, popped, counted, and returned. kequals the number of nodes. The answer is the largest value, found after a full traversal. For the second example tree withk = 6, the method returns 6.- A node with only a right child. In
[3, 1, 4, null, 2], node 1 has a right child 2 and no left one. After 1 is popped and counted,node = node.rightsends the walk into 2 before 3 is popped. The order 1, 2, 3, 4 comes out right. - Counting when pushing. A node pushed on the way down is not yet visited, because its left subtree comes first. Count when popping.
- Forgetting
node = node.right. After the pop,nodestill points at the popped node, so the inner loop would push it again and the walk would never move on. The right subtrees would never be seen. - Treating
kas a 0-based index.k = 1is the smallest value, not the second. - Sorting all values. Copying the values into a list and sorting works, but costs
O(n log n), and ignores the order the tree already has.
Testing
[3, 1, 4, null, 2],k = 1gives1.[5, 3, 6, 2, 4, null, null, 1],k = 3gives3.[5, 3, 6, 2, 4, null, null, 1],k = 6gives6.[5, 3, 6, 2, 4, null, null, 1],k = 1gives1.
