Problem overview
The function kthSmallest receives the root of a binary search tree (BST) and an integer k, counting from 1, and returns the k-th smallest value in the tree. The tree has 1 to 1,000 nodes with distinct values, and 1 <= k <= n.
For the tree with root 3, left child 1 (which has a right child 2) and right child 4, the values in increasing order are 1, 2, 3, 4, so k = 1 gives 1. For the tree with root 5, children 3 and 6, then 2 and 4 under 3, and 1 under 2, the sorted values are 1 to 6; k = 3 gives 3 and k = 6 gives 6.
Approach: in-order traversal that stops early
In a BST every value in a node's left subtree is smaller than the node, and every value in its right subtree is larger. An in-order traversal, which visits the left subtree, then the node, then the right subtree, therefore visits the values in increasing order. The k-th node that this traversal visits holds the answer.
You could collect all values into a slice and index it, but that always walks the whole tree. The follow-up asks you to stop as soon as the answer is known, in O(h + k) time, where h is the tree's height. With a recursive traversal, stopping early is awkward, so the lesson uses an explicit stack.
The stack holds the nodes you have passed on the way down and not yet visited. From any node, the smallest value in its part of the tree is at the leftmost node, so you push each node as you go left. When there is no more left, the top of the stack is the next node in order. You visit it, count it, and then handle its right subtree the same way.
Step by step
- Declare
stackas an empty slice of node pointers,[]*TreeNode{}.appendadds to the end and slicing off the last element removes it, so a slice works as a stack. - Declare
nodeand set it toroot. It is the root of the part of the tree still to be walked. - Declare
countat 0, the number of nodes visited so far. - Write the outer loop, which runs while
node != nilor the stack is not empty. The traversal is over only when no part of the tree is left to walk and no node waits on the stack. - Add an inner loop that runs while
node != nil, to go to the leftmost node. - Push
nodeonto the stack. It waits there until the smaller values in its left subtree have been visited. - Move to the left child:
node = node.Left. - Pop the top of the stack into
node. It is the smallest unvisited value, because everything smaller, including its left subtree, has been visited. - Count it:
count++. - If
count == k, returnnode.Val. The rest of the tree is never visited. - Otherwise move to the right child. The next values in order are in this node's right subtree, and they 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 always returns first, but Go requires the finalreturn.
The solution in Go
func kthSmallest(root *TreeNode, k int) int {
stack := []*TreeNode{}
node := root
count := 0
for node != nil || len(stack) > 0 {
for node != nil {
stack = append(stack, node)
node = node.Left
}
node = stack[len(stack)-1]
stack = stack[:len(stack)-1]
count++
if count == k {
return node.Val
}
node = node.Right
}
return -1
}
Go details worth a closer look:
- A slice as a stack.
stack = append(stack, node)pushes. The pop takesstack[len(stack)-1]and then re-slices withstack[:len(stack)-1], which shortens the slice without copying elements.appendreturns the slice, so its result must be assigned back. - A
forwith only a condition is Go's while loop. Both loops here use it. The outer condition,node != nil || len(stack) > 0, combines two cases with||. - Reusing
node. The same variable walks down, receives the popped node and then moves right. After the pop there is no need for another variable. - Element type.
[]*TreeNodeholds pointers, so pushing a node copies only the pointer, not the node. - The final
return -1is there only to satisfy the compiler, which cannot prove that the loop returns.
Complexity
Time: O(h + k). The traversal stops after k visits, which are k pops. Every node pushed is either popped by then or still on the stack, and the stack never holds more than h nodes, so there are at most k + h pushes. In the worst case, k = n, this becomes O(n).
Space: O(h). The stack holds the path from the root to the current node, which is at most the tree's height: about log n for a balanced tree and n for a chain.
Edge cases and common mistakes
- A single node. A root with value 9 and
k = 1gives 9: it is pushed, popped and counted. - The largest value.
k = nmakes the traversal visit every node. For a chain that only goes left, such as 5, 4, 3, 2, 1 inserted in that order,k = 5gives 5 after all five nodes are pushed. - Counting before the pop. Counting while pushing would count nodes in the order of the descent, not in sorted order, and give wrong answers.
- Forgetting the move to the right child. Without
node = node.Right, the next pass would pop a node of an earlier level and skip every right subtree. - Counting from 0.
kstarts at 1, so the check iscount == kafter incrementing, andk = 1returns the smallest value. - Sorting everything. Collecting all values and sorting ignores the BST order and costs
O(n log n).
Testing
- Root 3, with 1 on the left (and 2 as its right child) and 4 on the right,
k = 1gives 1. - Root 5, with 3 and 6 below it, 2 and 4 below 3, and 1 below 2,
k = 3gives 3. - The same tree with
k = 6gives 6. - The same tree with
k = 4gives 4. - A single node 9 with
k = 1gives 9.
