Binary Search Trees (BSTs) are powerful data structures that maintain a sorted order of elements. One common problem involving BSTs is finding the Kth smallest or largest element. In this article, we'll explore different approaches to solve this problem efficiently using Java.

Understanding the Problem

Given a BST and a positive integer K, our task is to find the Kth smallest or largest element in the tree. For example, in the following BST:

     5
   /   \
  3     7
 / \   / \
2   4 6   8

• The 3rd smallest element is 4
• The 2nd largest element is 7

Approach 1: In-order Traversal

The simplest approach leverages the fact that an in-order traversal of a BST visits nodes in ascending order.

For Kth Smallest Element:

class Solution {
    int count = 0;
    int result = -1;

    public int kthSmallest(TreeNode root, int k) {
        inorderTraversal(root, k);
        return result;
    }

    private void inorderTraversal(TreeNode node, int k) {
        if (node == null || count >= k) return;

        inorderTraversal(node.left, k);
        
        count++;
        if (count == k) {
            result = node.val;
            return;
        }

        inorderTraversal(node.right, k);
    }
}

For Kth Largest Element:

To find the Kth largest element, we can perform a reverse in-order traversal (right, root, left) and stop when we reach the Kth element.

class Solution {
    int count = 0;
    int result = -1;

    public int kthLargest(TreeNode root, int k) {
        reverseInorderTraversal(root, k);
        return result;
    }

    private void reverseInorderTraversal(TreeNode node, int k) {
        if (node == null || count >= k) return;

        reverseInorderTraversal(node.right, k);
        
        count++;
        if (count == k) {
            result = node.val;
            return;
        }

        reverseInorderTraversal(node.left, k);
    }
}

Time Complexity: O(N), where N is the number of nodes in the BST. Space Complexity: O(H) for the recursive call stack, where H is the height of the tree.

Approach 2: Augmented BST

For scenarios where we need to perform this operation frequently, we can augment the BST by storing the count of nodes in each subtree.

class AugmentedTreeNode {
    int val;
    int leftCount;
    AugmentedTreeNode left, right;

    AugmentedTreeNode(int val) {
        this.val = val;
        this.leftCount = 0;
    }
}

class AugmentedBST {
    AugmentedTreeNode root;

    public int kthSmallest(int k) {
        return kthSmallestHelper(root, k);
    }

    private int kthSmallestHelper(AugmentedTreeNode node, int k) {
        if (node == null) return -1;

        int leftCount = node.leftCount;

        if (k <= leftCount) {
            return kthSmallestHelper(node.left, k);
        } else if (k > leftCount + 1) {
            return kthSmallestHelper(node.right, k - leftCount - 1);
        } else {
            return node.val;
        }
    }

    // Insert method to maintain leftCount
    public void insert(int val) {
        root = insertHelper(root, val);
    }

    private AugmentedTreeNode insertHelper(AugmentedTreeNode node, int val) {
        if (node == null) return new AugmentedTreeNode(val);

        if (val < node.val) {
            node.left = insertHelper(node.left, val);
            node.leftCount++;
        } else {
            node.right = insertHelper(node.right, val);
        }

        return node;
    }
}

Time Complexity: O(H) for both insert and kthSmallest operations, where H is the height of the tree. Space Complexity: O(N) to store the additional count information.

When to Use Each Approach

1. Use the in-order traversal approach when:

   • The BST structure doesn't change frequently
   • Memory usage is a concern
   • You need a simple, straightforward solution

2. Use the augmented BST approach when:

   • You need to perform the Kth smallest/largest queries frequently
   • The BST is relatively stable (not many insertions/deletions)
   • You can afford the extra space for storing counts

Conclusion

Finding the Kth smallest or largest element in a BST is a common problem with multiple solutions. The in-order traversal method is simple and works well for occasional queries, while the augmented BST approach offers faster queries at the cost of additional space and maintenance during insertions and deletions.