Binary Search TreesProblem 15 of 22

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

Brute Force
Time: O(n log n)
Space: O(n)
Optimal
Time: O(n)
Space: O(1)

Problem Statement

Given a Binary Search Tree, find the median of the BST in O(n) time and O(1) space. The median is the middle element when elements are sorted. If the BST has even number of nodes, return the average of two middle elements.

Example:

6 / \ 3 8 / \ 1 4 Inorder: [1, 3, 4, 6, 8] Median: 4 (middle element, n=5) 6 / \ 3 8 / \ / \ 1 4 7 9 Inorder: [1, 3, 4, 6, 7, 8, 9] - wait, that's 7 nodes Actually: [1, 3, 4, 6, 7, 8, 9] Median: 6 (middle element, n=7)

Approach 1: Brute Force (Inorder to Array)

Intuition

Get sorted array via inorder traversal, then find the middle element(s).

Algorithm

  1. Perform inorder traversal to get sorted array
  2. If n is odd, return arr[n/2]
  3. If n is even, return (arr[n/2 - 1] + arr[n/2]) / 2
java
import java.util.*;

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

public class Solution {
    private void inorder(TreeNode root, List<Integer> nodes) {
        if (root == null) return;
        inorder(root.left, nodes);
        nodes.add(root.val);
        inorder(root.right, nodes);
    }
    
    public double findMedianBruteForce(TreeNode root) {
        List<Integer> nodes = new ArrayList<>();
        inorder(root, nodes);
        
        int n = nodes.size();
        if (n == 0) return 0;
        
        if (n % 2 == 1) {
            return nodes.get(n / 2);
        } else {
            return (nodes.get(n / 2 - 1) + nodes.get(n / 2)) / 2.0;
        }
    }
}

Complexity Analysis

  • Time Complexity: O(n) - Complete traversal
  • Space Complexity: O(n) - Storing all nodes

Approach 2: Optimal (Morris Traversal with Count)

Intuition

Use Morris Traversal for O(1) space. First count nodes, then find the median position during second Morris traversal.

Algorithm

  1. First pass: Count total nodes using Morris traversal
  2. Second pass: Find median element(s) at position(s) (n/2) and (n/2 + 1) for even, or (n/2 + 1) for odd
java
class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int x) { val = x; }
}

public class Solution {
    private int countNodes(TreeNode root) {
        int count = 0;
        TreeNode curr = root;
        
        while (curr != null) {
            if (curr.left == null) {
                count++;
                curr = curr.right;
            } else {
                TreeNode pred = curr.left;
                while (pred.right != null && pred.right != curr) {
                    pred = pred.right;
                }
                
                if (pred.right == null) {
                    pred.right = curr;
                    curr = curr.left;
                } else {
                    pred.right = null;
                    count++;
                    curr = curr.right;
                }
            }
        }
        
        return count;
    }
    
    public double findMedianOptimal(TreeNode root) {
        if (root == null) return 0;
        
        int n = countNodes(root);
        
        int count = 0;
        TreeNode curr = root;
        TreeNode first = null;
        TreeNode second = null;
        
        while (curr != null) {
            if (curr.left == null) {
                count++;
                
                if (n % 2 == 1 && count == (n + 1) / 2) {
                    second = curr;
                } else if (n % 2 == 0) {
                    if (count == n / 2) first = curr;
                    if (count == n / 2 + 1) second = curr;
                }
                
                curr = curr.right;
            } else {
                TreeNode pred = curr.left;
                while (pred.right != null && pred.right != curr) {
                    pred = pred.right;
                }
                
                if (pred.right == null) {
                    pred.right = curr;
                    curr = curr.left;
                } else {
                    pred.right = null;
                    count++;
                    
                    if (n % 2 == 1 && count == (n + 1) / 2) {
                        second = curr;
                    } else if (n % 2 == 0) {
                        if (count == n / 2) first = curr;
                        if (count == n / 2 + 1) second = curr;
                    }
                    
                    curr = curr.right;
                }
            }
        }
        
        if (n % 2 == 1) {
            return second.val;
        } else {
            return (first.val + second.val) / 2.0;
        }
    }
}

Complexity Analysis

  • Time Complexity: O(n) - Two Morris traversals
  • Space Complexity: O(1) - No extra space (Morris modifies and restores)

Approach 3: Optimized Single Pass (If Node Count Known)

Intuition

If we already know the node count (or have it stored), we can find median in a single Morris traversal.

Algorithm

  1. If count is precomputed, directly use it
  2. Single Morris traversal to find median position
java
class TreeNode {
    int val;
    int count;  // Augmented: number of nodes in subtree
    TreeNode left, right;
    TreeNode(int x) { val = x; count = 1; }
}

public class Solution {
    public double findMedianAugmented(TreeNode root, int n) {
        if (root == null || n == 0) return 0;
        
        int count = 0;
        TreeNode curr = root;
        TreeNode first = null;
        
        while (curr != null) {
            if (curr.left == null) {
                count++;
                
                if (n % 2 == 1 && count == (n + 1) / 2) {
                    return curr.val;
                }
                if (n % 2 == 0) {
                    if (count == n / 2) first = curr;
                    if (count == n / 2 + 1) {
                        return (first.val + curr.val) / 2.0;
                    }
                }
                
                curr = curr.right;
            } else {
                TreeNode pred = curr.left;
                while (pred.right != null && pred.right != curr) {
                    pred = pred.right;
                }
                
                if (pred.right == null) {
                    pred.right = curr;
                    curr = curr.left;
                } else {
                    pred.right = null;
                    count++;
                    
                    if (n % 2 == 1 && count == (n + 1) / 2) {
                        return curr.val;
                    }
                    if (n % 2 == 0) {
                        if (count == n / 2) first = curr;
                        if (count == n / 2 + 1) {
                            return (first.val + curr.val) / 2.0;
                        }
                    }
                    
                    curr = curr.right;
                }
            }
        }
        
        return 0;
    }
}

Complexity Analysis

  • Time Complexity: O(n) - Single Morris traversal
  • Space Complexity: O(1)

Key Takeaways

  1. Morris Traversal: Achieves O(1) space by temporarily modifying tree
  2. Two passes needed: First to count, second to find median (unless count is known)
  3. Handle even/odd: Average of two middle elements for even count
  4. Inorder property: BST inorder is sorted, making median finding straightforward
  5. Morris traversal is perfect for space-constrained scenarios
  6. The tree is restored after Morris traversal
  7. Can early terminate once median is found (optimization)