Find median of BST in O(n) time and O(1) space in Python

Finding the median of a Binary Search Tree (BST) efficiently requires understanding that an inorder traversal of a BST gives nodes in sorted order. The median is the middle element for odd number of nodes, or the average of two middle elements for even number of nodes.

For a BST with n nodes:

• Odd nodes: median = (n+1)/2th node

• Even nodes: median = average of (n/2)th and (n/2+1)th nodes

We'll use Morris Traversal to achieve O(n) time complexity with O(1) space complexity by avoiding recursion stack.

Example Tree

Algorithm Steps

The solution uses Morris Traversal in two phases:

1. Count nodes: Traverse the BST to get total node count

2. Find median: Traverse again to find the median element(s)

Implementation

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

def count_nodes(root):
    """Count total nodes using Morris Traversal"""
    if root is None:
        return 0
    
    node_count = 0
    current = root
    
    while current is not None:
        if current.left is None:
            node_count += 1
            current = current.right
        else:
            # Find inorder predecessor
            predecessor = current.left
            while predecessor.right is not None and predecessor.right != current:
                predecessor = predecessor.right
            
            if predecessor.right is None:
                # Create temporary link
                predecessor.right = current
                current = current.left
            else:
                # Remove temporary link and count node
                predecessor.right = None
                node_count += 1
                current = current.right
    
    return node_count

def find_median(root):
    """Find median using Morris Traversal with O(1) space"""
    if root is None:
        return 0
    
    node_count = count_nodes(root)
    current = root
    count = 0
    prev_data = 0
    
    while current is not None:
        if current.left is None:
            count += 1
            
            # Check if we found the median
            if node_count % 2 == 1 and count == (node_count + 1) // 2:
                return current.data
            elif node_count % 2 == 0 and count == node_count // 2 + 1:
                return (prev_data + current.data) / 2
            
            prev_data = current.data
            current = current.right
        else:
            # Find inorder predecessor
            predecessor = current.left
            while predecessor.right is not None and predecessor.right != current:
                predecessor = predecessor.right
            
            if predecessor.right is None:
                # Create temporary link
                predecessor.right = current
                current = current.left
            else:
                # Remove temporary link
                predecessor.right = None
                count += 1
                
                # Check if we found the median
                if node_count % 2 == 1 and count == (node_count + 1) // 2:
                    return current.data
                elif node_count % 2 == 0 and count == node_count // 2 + 1:
                    return (prev_data + current.data) / 2
                
                prev_data = current.data
                current = current.right
    
    return 0

# Create the BST from the example
root = TreeNode(7)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(2)
root.left.right = TreeNode(5)
root.right.left = TreeNode(8)
root.right.right = TreeNode(10)

print("Median of BST:", find_median(root))

Median of BST: 7

How Morris Traversal Works

Morris Traversal uses temporary links to simulate recursion without using extra space:

1. If left child is None, process current node and move right

2. If left child exists, find the rightmost node in left subtree (predecessor)

3. If predecessor's right is None, create temporary link and go left

4. If predecessor's right points to current, remove link, process node, and go right

Time and Space Complexity

Conclusion

Morris Traversal provides an elegant solution to find BST median in O(n) time with O(1) space. It temporarily modifies the tree structure during traversal but restores it completely, making it space-efficient for large datasets.