medium Kth Smallest Element in a BST

Problem Statement

Given the root of a binary search tree (BST) and an integer k, return the k-th smallest value in the tree.

The values in the BST are unique.

. Examples

Example 1:

• Input: root = [5, 3, 8, 2, 4, 7, 9, 1], k = 3

      5
     / \
    3   8
   / \  / \
  2  4 7   9
 / 
 1

• Expected Output: 3
• Justification: The sorted order of the elements is [1, 2, 3, 4, 5, 7, 8, 9]. The 3rd smallest value is 3.

Example 2:

• Input: root = [6, 4, 8, 2, 5, 7, 9, 1, 3], k = 5

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

• Expected Output: 5
• Justification: The sorted order of the elements is [1, 2, 3, 4, 5, 6, 7, 8, 9]. The 5th smallest value is 5.

Example 3:

• Input: root = [10, 5, 15, 3, 7, 13, 18, 2, 4, 6, 8], k = 6

      10
     /  \
    5    15
   /|    / \
  3 7   13 18
 /| |   |
2 4 6   8

• Expected Output: 7
• Justification: The sorted order of the elements is [2, 3, 4, 5, 6, 7, 8, 10, 13, 15, 18]. The 6th smallest value is 7.

Constraints:

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

Solution

To solve this problem, we can use an in-order traversal of the binary search tree (BST). In-order traversal of a BST gives us the elements in sorted order. By traversing the tree in this way, we can count the nodes as we go and stop when we reach the k-th smallest node.

This approach works because the properties of the BST guarantee that an in-order traversal visits nodes in ascending order. This method is efficient and ensures that we visit the minimum number of nodes necessary to find the k-th smallest element.

Step-by-Step Algorithm

1. Initialize a stack to help with the in-order traversal.
2. Set a pointer (current) to the root node of the BST.
3. Initialize a counter to keep track of the number of nodes visited.
4. While the stack is not empty or the current node is not null:
   • While the current node is not null:
     • Push the current node onto the stack.
     • Move to the left child of the current node.
   • Pop a node from the stack and set it as the current node.
   • Increment the counter by 1.
   • If the counter equals k, return the value of the current node (current.val).
   • Move to the right child of the current node.
5. Return -1 if the tree has fewer than k nodes (this condition is not necessary here as problem guarantees valid k).

Algorithm Walkthrough

Given the input:

• root: [6, 4, 8, 2, 5, 7, 9, 1, 3]
• k: 5

Step-by-Step Walkthrough:

1. Initialize:

   • stack = []
   • current = root (6)
   • count = 0

2. First outer while loop:

   • current = 6
   • stack = []
   • Move to left child:
     • Push 6 to stack: stack = [6]
     • current = 4
   • Move to left child:
     • Push 4 to stack: stack = [6, 4]
     • current = 2
   • Move to left child:
     • Push 2 to stack: stack = [6, 4, 2]
     • current = 1
   • Move to left child:
     • Push 1 to stack: stack = [6, 4, 2, 1]
     • current = null (left child of 1)

3. First pop from stack:

   • current = stack.pop() = 1
   • Increment count = 1
   • count is not equal to k
   • Move to right child:
     • current = null (right child of 1)

4. Second pop from stack:

   • current = stack.pop() = 2
   • Increment count = 2
   • count is not equal to k
   • Move to right child:
     • current = 3

5. Second outer while loop:

   • current = 3
   • stack = [6, 4]
   • Move to left child:
     • current = null (left child of 3)

6. Third pop from stack:

   • current = stack.pop() = 3
   • Increment count = 3
   • count is not equal to k
   • Move to right child:
     • current = null (right child of 3)

7. Fourth pop from stack:

   • current = stack.pop() = 4
   • Increment count = 4
   • count is not equal to k
   • Move to right child:
     • current = 5

8. Third outer while loop:

   • current = 5
   • stack = [6]
   • Move to left child:
     • current = null (left child of 5)

9. Fifth pop from stack:

   • current = stack.pop() = 5
   • Increment count = 5
   • count is equal to k
   • Return current.val = 5

At this point, we have found the k-th smallest element which is 5.

Code

import java.util.Stack;

// class TreeNode {
//     int val;
//     TreeNode left;
//     TreeNode right;
//     TreeNode(int x) { val = x; }
// }

class Solution {

  public int kthSmallest(TreeNode root, int k) {
    Stack<TreeNode> stack = new Stack<>();
    TreeNode current = root;
    int count = 0;

    // Traverse the tree in in-order fashion
    while (current != null || !stack.isEmpty()) {
      // Reach the leftmost node of the current node
      while (current != null) {
        stack.push(current);
        current = current.left;
      }
      // Current must be null at this point
      current = stack.pop();
      count++;
      // If count equals k, we've found the k-th smallest element
      if (count == k) return current.val;
      // Move to the right node
      current = current.right;
    }

    return -1; // Return -1 if k is out of bounds
  }

  public static void main(String[] args) {
    Solution sol = new Solution();

    // Example 1
    TreeNode root1 = new TreeNode(5);
    root1.left = new TreeNode(3);
    root1.left.left = new TreeNode(2);
    root1.left.left.left = new TreeNode(1);
    root1.left.right = new TreeNode(4);
    root1.right = new TreeNode(8);
    root1.right.left = new TreeNode(7);
    root1.right.right = new TreeNode(9);
    System.out.println(sol.kthSmallest(root1, 3)); // Output: 3

    // Example 2
    TreeNode root2 = new TreeNode(6);
    root2.left = new TreeNode(4);
    root2.left.left = new TreeNode(2);
    root2.left.left.left = new TreeNode(1);
    root2.left.left.right = new TreeNode(3);
    root2.left.right = new TreeNode(5);
    root2.right = new TreeNode(8);
    root2.right.left = new TreeNode(7);
    root2.right.right = new TreeNode(9);
    System.out.println(sol.kthSmallest(root2, 5)); // Output: 5

    // Example 3
    TreeNode root3 = new TreeNode(10);
    root3.left = new TreeNode(5);
    root3.left.left = new TreeNode(3);
    root3.left.left.left = new TreeNode(2);
    root3.left.left.right = new TreeNode(4);
    root3.left.right = new TreeNode(7);
    root3.left.right.left = new TreeNode(6);
    root3.left.right.right = new TreeNode(8);
    root3.right = new TreeNode(15);
    root3.right.left = new TreeNode(13);
    root3.right.right = new TreeNode(18);
    System.out.println(sol.kthSmallest(root3, 6)); // Output: 7
  }
}

Complexity Analysis

Time Complexity

The time complexity of the algorithm is , where (H) is the height of the tree. In the worst case, (H) could be equal to (N) (if the tree is skewed), leading to a time complexity of . However, for a balanced tree, (H = log N), and the time complexity becomes .

Space Complexity

The space complexity is , which accounts for the stack used during the in-order traversal. In the worst case, this could also be for a skewed tree. For a balanced tree, the space complexity is .