Program to find number of nodes in a range in Python

Given a Binary Search Tree (BST) and two bounds l and r, we need to count all nodes whose values fall within the range [l, r] (inclusive).

For example, if we have a BST with nodes and l = 7, r = 13, we count nodes with values 8, 10, and 12, giving us a result of 3.

Algorithm

We use an iterative approach with a stack to traverse the BST efficiently −

• Initialize a stack with the root node and a counter

• While the stack is not empty:

  • Pop a node from the stack

  • If the node's value is in range [l, r], increment counter and add both children

  • If the node's value is less than l, only explore the right subtree

  • If the node's value is greater than r, only explore the left subtree

Implementation

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

class Solution:
    def solve(self, root, l, r):
        stack, count = [root], 0
        
        while stack:
            node = stack.pop()
            if node:
                if l <= node.data <= r:
                    count += 1
                    stack += [node.right, node.left]
                elif node.data < l:
                    stack += [node.right]
                else:
                    stack += [node.left]
        
        return count

# Create the BST
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)

# Test the solution
ob = Solution()
result = ob.solve(root, 7, 13)
print(result)

3

How It Works

The algorithm leverages the BST property to optimize traversal. When a node's value is outside our target range, we can skip entire subtrees:

• If node.data < l: All nodes in the left subtree are also less than l, so we only check the right subtree

• If node.data > r: All nodes in the right subtree are also greater than r, so we only check the left subtree

• If l <= node.data <= r: The node is in range, so we count it and check both subtrees

Time Complexity

The time complexity is O(n) in the worst case, where n is the number of nodes. However, due to the BST property, we often skip many nodes, making it more efficient in practice.

Conclusion

This iterative stack-based approach efficiently counts nodes in a BST within a given range by leveraging BST properties to prune unnecessary subtrees. The algorithm is both space-efficient and handles the range checking optimally.