Sum of k smallest elements in BST

Problem

Given a Binary Search Tree (BST), the task is to find the sum of all elements that are smaller than or equal to the Kth smallest element in the BST.

Examples

Example 1

Input:  
K = 3  
Tree structure:  
		   7  
		 /   \
		6     9  
	  /      /   \  
	 1     8     11  

Output:  
14  

Explanation:  
The Kth smallest element is the 3rd smallest, i.e., `7`.  
The sum of values smaller than or equal to `7` in the BST is `1 + 6 + 7 = 14`.  

Solution

Method 1 - Inorder Traversal

1. The in-order traversal method will visit nodes in sorted order due to the BST property.
2. The Kth smallest element relates to the position of nodes visited during in-order traversal.
3. Once the Kth smallest value is identified, we can sum up all nodes visited in traversal until we reach this value.

Approach

1. Perform an in-order traversal of the BST:
   • Keep track of nodes visited in sorted order.
   • Maintain a running sum of node values during traversal.
2. Stop traversal when we find the Kth smallest element (based on the node count during traversal).
3. The running sum up to the Kth value is the result.

Code

Java

class Solution {

    private int nodeCount = 0; // Counter to track number of visited nodes
    private int sum = 0; // Sum of all nodes visited

    // Function to find the sum of elements <= Kth smallest in BST
    public int sumOfSmallerOrEqualToK(TreeNode root, int K) {
        inOrderTraversal(root, K);
        return sum;
    }

    // Helper function for in-order traversal
    private void inOrderTraversal(TreeNode root, int K) {
        if (root == null) {
            return;
        }

        // Traverse left subtree
        inOrderTraversal(root.left, K);

        // Process current node
        nodeCount++;
        if (nodeCount <= K) {
            sum += root.val;
        } else {
            return; // Stop traversal once Kth smallest is reached
        }

        // Traverse right subtree
        inOrderTraversal(root.right, K);
    }
}

Python

class Solution:
    def __init__(self):
        self.node_count = 0  # Counter to track number of visited nodes
        self.sum = 0  # Sum of all nodes visited
    
    def sum_of_smaller_or_equal_to_k(self, root, K):
        self.in_order_traversal(root, K)
        return self.sum
    
    # Helper function for in-order traversal
    def in_order_traversal(self, root, K):
        if not root:
            return
        
        # Traverse left subtree
        self.in_order_traversal(root.left, K)
        
        # Process current node
        self.node_count += 1
        if self.node_count <= K:
            self.sum += root.val
        else:
            return  # Stop traversal once Kth smallest is reached
        
        # Traverse right subtree
        self.in_order_traversal(root.right, K)

Complexity

• ⏰ Time complexity: O(n), since we may need to traverse all nodes in the BST.
• 🧺 Space complexity: O(h + k). For recursion (in-order traversal): O(H) (where h is the height of the BST).
  • For storing values during traversal (array-based approach): O(k) for intermediate storage.