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. The values are all different, and k is between 1 and the number of nodes.
For the tree 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, so k = 1 gives 1. For the tree with root 5, children 3 and 6, and below 3 the nodes 2 and 4 (with 1 under 2), the order is 1, 2, 3, 4, 5, 6: k = 3 gives 3, and k = 6 gives 6.
Approach: in-order traversal with a stack
An in-order traversal visits the left subtree, then the node, then the right subtree. In a binary search tree, that visits the values in increasing order, so the k-th node visited holds the answer.
A recursive traversal is simple but runs to the end of the tree. The follow-up asks for O(h + k) time, which means stopping at the k-th node. That is easier with an explicit stack, where you can just return from the middle of the loop.
The stack holds nodes you walked past on the way down and haven't visited yet. The procedure is:
- From the current node, go left as far as possible, pushing every node. The smallest value of that part of the tree is at its leftmost node.
- Pop a node. It is the smallest value not yet visited, because all smaller values, including its entire left subtree, were visited earlier. Count it.
- If the count equals
k, return its value. Otherwise continue with its right subtree.
On the second example with k = 3: the descent pushes 5, 3, 2 and 1. Pop 1: count 1, no right child. Pop 2: count 2, no right child. Pop 3: count 3, which equals k, so the answer is 3, and the nodes 4, 5 and 6 are never touched.
Step by step
- Create
stackas an empty[TreeNode]withvar.appendpushes andremoveLast()pops, so an array works as a stack. - Declare
nodewithvarasroot. It is the root of the part of the tree still to be walked. When it isnil, the next node to visit is on top of the stack. - Declare
countat 0: the number of nodes visited so far. - Loop while
node != nil || !stack.isEmpty. The traversal is over only when nothing is left to descend into and nothing waits on the stack. - Descend left with an inner
while let current = node. - Push
current. Its left subtree holds smaller values, so it waits. - Set
nodetocurrent.left. - After the inner loop, pop the last node into a new constant
currentwithremoveLast(). It is the smallest unvisited value. - Visit it:
count += 1. - Check
count == k. If so, returncurrent.val. The rest of the tree is never visited. - Otherwise go right: set
nodetocurrent.right, so the next pass walks through that subtree. Its values are bigger thancurrentbut smaller than the nodes still on the stack. - Return -1 after the loop. Since
kis at most the number of nodes, this line never runs, but Swift requires a return on every path.
The solution in Swift
class Solution {
func kthSmallest(_ root: TreeNode?, _ k: Int) -> Int {
var stack: [TreeNode] = []
var node = root
var count = 0
while node != nil || !stack.isEmpty {
while let current = node {
stack.append(current)
node = current.left
}
let current = stack.removeLast()
count += 1
if count == k {
return current.val
}
node = current.right
}
return -1
}
}
Swift details worth knowing:
- Two variables called
current. The one inwhile let current = nodeexists only inside that inner loop. After the loop ends,let current = stack.removeLast()declares a new constant in the outer loop's scope, which is why the same name can be reused. removeLast()returns the element. It both removes the last item and gives it back, and it crashes on an empty array. Here the stack can't be empty at that point: ifnodeisniland the loop condition holds, the stack has something in it, and ifnodewasn'tnil, the inner loop pushed at least one node.- Arrays of class instances.
[TreeNode]stores references, so pushing a node doesn't copy it. - Optionals in a condition.
node != nilcompares an optional withnilwithout unwrapping it, and the stack check uses theisEmptyproperty, which is clearer thanstack.count == 0.
Complexity
Time: O(h + k), where h is the height of the tree. The first descent costs up to h pushes, and after that each of the k visits pops one node and pushes the nodes of a left chain. Every node is pushed and popped at most once, so the walk never exceeds O(n).
Space: O(h). The stack holds the path from the root down to the current node, at most h nodes.
Edge cases and common mistakes
- A single node. The tree with only 7 and
k = 1pushes the node, pops it, counts 1 and returns 7. k = 1. The answer is the leftmost node, found after the first descent.kequal to the number of nodes. The walk visits every node and returns the largest value, the rightmost one.- A skewed tree. For a chain going only left, the first descent pushes every node (the stack depth is
n), and then the pops run in increasing order. - Counting when pushing. Pushes are not visits. Count a node when it is popped, in sorted order, not when you first reach it.
- Forgetting to go right. After a visit, the node's right subtree contains the next larger values. Without
node = current.rightthey are skipped. - Counting from 0.
kstarts at 1, andcountstarts at 0, so comparing after incrementing is right.
Testing
- Root 3 with children 1 and 4 and 2 under 1,
k = 1gives 1. - Root 5 with children 3 and 6, then 2 and 4 under 3 and 1 under 2:
k = 3gives 3,k = 1gives 1 andk = 6gives 6. - For the same tree,
k = 2gives 2,k = 4gives 4 andk = 5gives 5. - A single node holding 7,
k = 1, gives 7.
