- Is
kalways valid? Within the range of 1 to the size of the BST?
Input:
5
/ \
3 6
/ \ \
2 4 7
k = 3 / 5
Output: 4 / 6
- The inorder traversal of BST is in the ascending order, so we can traverse the BST in inorder and increment the index until it reaches
k.
// Iterative
fun kthSmallest(root: TreeNode?, k: Int): Int {
if (root == null) -1
val stack = Stack<TreeNode>()
var index = 0
var node: TreeNode? = root
while (!stack.isEmpty() || node != null) {
while (node != null) {
stack.push(node)
node = node.left
}
node = stack.pop()
index++
if (index == k) return node.`val`
node = node.right
}
return -1
}
// Recursive
private var i = 0
private var result = -1
fun kthSmallest(root: TreeNode?, k: Int): Int {
inorder(root, k)
return result
}
private fun inorder(root: TreeNode?, k: Int) {
if (root == null) return
inorder(root.left, k)
i++
if (i == k) {
result = root.`val`
return
}
inorder(root.right, k)
}- Time Complexity:
O(h + k),hfor stack and k-th elements. - Space Complexty:
O(h)for stack.
- We can count the number of nodes in the left subtree, if the count is greater than
k, we can search in the left subtree, otherwise we can search in the right subtree. We return when the count is equal tok.
fun kthSmallest(root: TreeNode?, k: Int): Int {
if (root == null) return -1
val leftCount = count(root.left)
// Minus 1 because the root is counted.
if (leftCount < k - 1) {
// Minus leftCount because we have already counted the left subtree.
// In the right subtree, we search for the k - leftCount - 1-th element.
// See below example
return kthSmallest(root.right, k - leftCount - 1)
} else if (leftCount > k - 1) {
return kthSmallest(root.left, k)
} else {
return root.`val`
}
}
private fun count(root: TreeNode?): Int {
if (root == null) return 0
val leftCount = count(root.left)
val rightCount = count(root.right)
return leftCount + rightCount + 1
}- Time Complexity:
O(n^2)if we don't precompute the count, orO(n)if we precompute and preserve the count of nodes in each node. - Space Complexity:
O(h).
Why do we search for k - leftCount - 1-th element in the right subtree?
k = 5
leftCount(5) = 4
5 // start from 5 to search 5th smallest element
/ \
3 7
/ \ / \
2 4 6 8
/
1
k = 6
5 // start from 5 to search 6th smallest element
/ \
3 7 // start from 7 to search 6 - 4 - 1 = 1th smallest element
/ \ / \
2 4 6 8
/
1