Problem overview
The function kthSmallest receives the root of a binary search tree and an integer k, counted from 1, and returns the k-th smallest value in the tree. 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. All values are different, and k is between 1 and the number of nodes.
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. With k = 1 the answer is 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, so k = 3 gives 3 and k = 6 gives 6.
Approach: in-order traversal, stopped early
An in-order traversal visits the left subtree, then the node, then the right subtree. On a BST this visits the values in increasing order. So the k-th node the traversal visits holds the k-th smallest value, and the traversal can stop right there.
A recursive traversal is the textbook version, but stopping in the middle of recursion needs extra bookkeeping. The lesson uses a stack instead, which makes the order explicit. The stack holds the nodes passed on the way down that are still waiting to be visited. The idea in three moves:
- From the current node, go left as far as possible, pushing every node on the way. The last node reached has the smallest value in that part of the tree.
- Pop the top of the stack. Everything smaller than it has already been visited. Visit it: count it, and return its value if the count equals
k. - Otherwise, continue from its right child, whose values come next in the order.
For the second example tree and k = 3: the walk pushes 5, 3, 2 and 1, and pops 1 (count 1). Node 1 has no right child, so the next pop is 2 (count 2), then 3 (count 3), which is the answer.
Step by step
- Create
stackwithvalas an emptyArrayDeque<TreeNode>().addLastandremoveLastboth work at the end, so it acts as a stack. - Declare
nodewithvarand set it toroot. It is the root of the part of the tree still to walk. Its type isTreeNode?. - Declare
countwithvar, starting at 0: the number of nodes visited so far. - Write the outer loop that runs while
node != nullorstack.isNotEmpty(). The traversal is over only when there is no part of the tree left to walk and no node waiting on the stack. - Write an inner loop that runs while
node != null. It heads to the leftmost node, because each step left reaches smaller values. - Push
nodewithaddLast. It waits on the stack until its left subtree has been visited. - Move
nodetonode.left. - After the inner loop, pop the top with
removeLast()and store it innode. This is the smallest value not yet visited. After this line, Kotlin knowsnodeis notnull. - Visit it:
count++. - If
count == k, returnnode.value. The rest of the tree is never visited. - Otherwise move
nodetonode.right. The next pass walks that subtree, whose values are larger than this node but smaller than the nodes still on the stack. - Return -1 after the outer loop. It never runs, since
kis at most the number of nodes, but a function that returns anInthas to return one on every path.
The solution in Kotlin
class Solution {
fun kthSmallest(root: TreeNode?, k: Int): Int {
val stack = ArrayDeque<TreeNode>()
var node = root
var count = 0
while (node != null || stack.isNotEmpty()) {
while (node != null) {
stack.addLast(node)
node = node.left
}
node = stack.removeLast()
count++
if (count == k) {
return node.value
}
node = node.right
}
return -1
}
}
The Kotlin details worth a closer look:
- A stack of non-null nodes.
ArrayDeque<TreeNode>can't holdnull.addLast(node)compiles inside the inner loop becausenodeis smart-cast toTreeNodethere. - Assigning from
removeLast().nodehas the typeTreeNode?, andremoveLast()returns aTreeNode. After the assignment, the compiler knowsnodeis notnulluntil it is assigned again, sonode.valueandnode.rightare allowed. removeLast()on an empty deque. It throws an exception. It is safe here: the outer loop runs only ifnode != nullor the stack is not empty, and after the inner loopnodeisnull, so the stack must hold something.count++. It is used as a statement on its own line, and it increments thevar.- Values and
k. The tree's field is calledvalue, so there is no keyword clash, unlike a field namedval.
Complexity
Let h be the height of the tree.
Time: O(h + k). Reaching the first node takes h pushes. After that, every visit pops one node, and every node is pushed once, so reaching the k-th visit takes at most h + k pushes and k pops. In the worst case, k is the number of nodes, and the time is O(n).
Space: O(h). The stack never holds more than one root-to-leaf path. A balanced tree has h about log n, and a tree shaped like a list has h equal to n.
Edge cases and common mistakes
k = 1. The answer is the leftmost node. In the first example tree, the walk pushes 3 and 1, then pops 1 and returns it.kequals the number of nodes. The answer is the largest value. For the second example tree andk = 6it is 6.- A single node. The only valid
kis 1, and the first pop returns the root's value. - An off-by-one count.
kcounts from 1, andcountis incremented before it is compared withk, so the first visited node hascount == 1. Comparing before the increment would return the node after the right one. - Visiting before going left. That gives a preorder walk, whose order is not sorted, so the count no longer matches the ranks.
- Forgetting the right child. Without
node = node.right, the walk would never visit the values larger than a popped node unless they were already on the stack. - Stopping the outer loop on
node != nullalone. After the first pop,nodecan benullbecause the popped node had no right child, while the stack still holds the nodes whose turn comes next. The loop must also run while the stack is not empty.
Testing
- Root 3, left 1 (right child 2), right 4, with
k = 1gives1;k = 2gives2;k = 4gives4. - Root 5, children 3 and 6, then 2 and 4 under 3, and 1 under 2, with
k = 3gives3;k = 5gives5;k = 6gives6.
