medium Kth Smallest Element in a BST

Problem Statement

Given a root node of the Binary Search Tree (BST) and integer 'k'. Return the Kth smallest element among all node values of the binary tree.

Examples:

1. Example 1:
   Input:

       8
      / \
     3   10
    / \    \
   1   6    14
      /  \  /
     4   7  13
   

   k = 4
   Expected Output: 6
   Justification: The in-order traversal of the tree is [1, 3, 4, 6, 7, 8, 10, 13, 14]. The 4th element in this sequence is 6.

2. Example 2:
   Input:

       5
      / \
     2   6
    /
   1
   

   k = 3
   Expected Output: 5
   Justification: The in-order traversal of the tree is [1, 2, 5, 6]. The 3rd element in this sequence is 5.

3. Example 3:
   Input:

   1
    \
     3
    /
   2
   

   k = 2
   Expected Output: 2
   Justification: The in-order traversal of the tree is [1, 2, 3]. The 2nd element in this sequence is 2.

Constraints:

• The number of nodes in the tree is n.
• 1 <= k <= n <= 10⁴
• 0 <= Node.val <= 10⁴

Solution

The most intuitive approach to this problem is to perform an in-order traversal on the BST. Given the nature of BSTs, an in-order traversal will result in a sequence of nodes in ascending order. Thus, by collecting the nodes in a list during traversal, the kth element of this list would be our answer.

However, there is a more efficient approach. We don't necessarily need to traverse the entire BST. We can halt our in-order traversal as soon as we've visited the kth node. This is achieved by using a counter that's incremented every time we visit a node during our traversal. When the counter matches 'k', we've found our kth smallest element.

Algorithm Walkthrough:
Given the tree from Example 1:

       8
      / \
     3   10
    / \    \
   1   6    14
      /  \  /
     4   7  13

With k = 4, the steps are:

• Start at the root (8). Since there's a left child, move to the left child (3).
• From node 3, again move to its left child (1). It's the leftmost node, so count it as the first smallest element.
• Move up to node 3. It's the second smallest.
• Now move to its right child (6). Since 6 also has a left child, move to that (4). It's the third smallest.
• Move up to node 6. It's the fourth smallest - which is what we are looking for. Hence, we stop and return 6.

Here is the visual representation of the algorithm:

Code

Here is the code for this algorithm:

// class TreeNode {
//     int val;                // Value of the node.
//     TreeNode left;          // Reference to the left child.
//     TreeNode right;         // Reference to the right child.

//     // Constructor to initialize the node with its value.
//     TreeNode(int x) { val = x; }
// }

public class Solution {

  // `count` keeps track of the number of nodes we've traversed in-order.
  private int count = 0;

  // `result` will hold our final answer.
  private int result = 0;

  // This method is the public API that finds the kth smallest element in a BST.
  public int kthSmallest(TreeNode root, int k) {
    // Start the in-order traversal.
    traverse(root, k);

    // Once traversal is done, the `result` will hold our answer.
    return result;
  }

  // A recursive function to do an in-order traversal of the BST.
  // We stop traversing once we've visited `k` nodes.
  private void traverse(TreeNode node, int k) {
    // If the current node is null or we've already traversed k nodes, return.
    if (node == null || count >= k) return;

    // First, traverse the left subtree.
    traverse(node.left, k);

    // Increment the counter for the current node.
    count++;

    // If we've traversed exactly k nodes, this is our result.
    if (count == k) result = node.val;

    // Finally, traverse the right subtree.
    traverse(node.right, k);
  }

  public static void main(String[] args) {
    // Constructing the tree for testing.
    TreeNode example1 = new TreeNode(8);
    example1.left = new TreeNode(3);
    example1.left.left = new TreeNode(1);
    example1.left.right = new TreeNode(6);
    example1.left.right.left = new TreeNode(4);
    example1.left.right.right = new TreeNode(7);
    example1.right = new TreeNode(10);
    example1.right.right = new TreeNode(14);
    example1.right.right.left = new TreeNode(13);

    Solution solution = new Solution();

    // Test the kthSmallest method.
    System.out.println(solution.kthSmallest(example1, 4)); // Expected output: 6
  }
}

Complexity Analysis

The worst-case time complexity is (O(N)) where (N) is the number of nodes, in case we need to visit all nodes. However, on average, we only have to visit k nodes, making it O(k). The space complexity is (O(1)) if we disregard the space used by the recursion stack, otherwise, it's (O(h)) where (h) is the height of the tree.