Kth Smallest Element in a BST

Problem Description

Given the root of a binary search tree (BST), and an integer k, return the kth smallest value (1-indexed) of all the values of the nodes in the tree.

Brute Force Solution

A brute force approach involves traversing the BST, storing all the node values in an array, sorting the array, and then returning the element at index k-1. This leverages the fact that we can extract all the node values, and then just sort them into ascending order.

Algorithm

1. Perform an in-order traversal of the BST to extract all node values.
2. Store the node values in an array.
3. Sort the array in ascending order.
4. Return the element at index k-1.

Code

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def kthSmallest_bruteforce(root: TreeNode, k: int) -> int:
    values = []

    def inorder_traversal(node):
        if node:
            inorder_traversal(node.left)
            values.append(node.val)
            inorder_traversal(node.right)

    inorder_traversal(root)
    values.sort()
    return values[k - 1]

Time and Space Complexity

• Time Complexity: O(N log N), where N is the number of nodes in the BST. O(N) for traversal and O(N log N) for sorting.
• Space Complexity: O(N) to store the node values in the array.

Optimal Solution

An optimal solution leverages the properties of a BST and in-order traversal to find the kth smallest element without sorting. In an in-order traversal, the left subtree is visited first, then the current node, and then the right subtree. For a BST, this ensures that the nodes are visited in ascending order.

Algorithm

1. Perform an in-order traversal of the BST.
2. Keep track of the number of nodes visited.
3. When the number of nodes visited equals k, return the current node's value.

Code

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def kthSmallest(root: TreeNode, k: int) -> int:
    count = 0
    result = None

    def inorder_traversal(node):
        nonlocal count, result
        if not node:
            return

        inorder_traversal(node.left)
        count += 1
        if count == k:
            result = node.val
            return
        inorder_traversal(node.right)

    inorder_traversal(root)
    return result

Time and Space Complexity

• Time Complexity: O(H + k), where H is the height of the BST. In the worst case, H can be N (skewed tree), and in the best case, H can be log N (balanced tree). Therefore, in the average case, it will be closer to O(log N + k) and in the worst case O(N).
• Space Complexity: O(H) due to the recursive call stack, where H is the height of the tree. In the worst case, this is O(N), and in the best case, this is O(log N).

Edge Cases

• If the BST is empty (root is None), the problem statement guarantees that k is within the range of node count. In other words, the BST will not be empty.
• If k is 1, return the smallest element in the BST (leftmost node).
• If k is equal to the number of nodes, return the largest element in the BST (rightmost node).
• If k is out of range, the code, as written, will return None. However, problem constraints specify that 1 <= k <= n, so this case should not occur.

Follow-up: Optimizations for frequent modifications

If the BST is modified often with insert and delete operations, and we need to find the kth smallest element frequently, we can augment the BST by storing the size of each subtree in the node. This allows us to find the kth smallest element in O(log N) time, where N is the number of nodes in the tree.

When a node is inserted or deleted, we need to update the size of all the subtrees that contain the modified node. This update takes O(log N) time as well.

With subtree sizes available, finding the kth smallest element becomes:

1. Let size be the size of the left subtree of the current node.
2. If k == size + 1, then the current node is the kth smallest.
3. If k <= size, then the kth smallest element is in the left subtree.
4. If k > size + 1, then the kth smallest element is in the right subtree. We need to search for the (k - size - 1)th smallest element in the right subtree.

This approach maintains a time complexity of O(log N) for finding the kth smallest element, making it efficient for frequent modifications and queries.