230.kth-smallest-element-in-a-bst.md

230. Kth Smallest Element in a BST

Clarification Questions

• Is k always valid? Within the range of 1 to the size of the BST?

Test Cases

Normal Cases

Input: 
        5
      /   \
     3     6
    / \     \
   2   4     7
k = 3
Output: 4 

Edge / Corner Cases

Input: 
Output: 

Inorder Traversal

• 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), h for stack and k-th elements.
• Space Complexty: O(h) for stack.

Counting

• 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 to k.

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, or O(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