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

Suppose we have Binary Search Tree(BST), we have to find median of it. We know for even number of nodes, median = ((n/2th node + (n+1)/2th node) /2 For odd number of nodes, median = (n+1)/2th node.

So, if the input is like

then the output will be 7

To solve this, we will follow these steps −

• if root is same as None, then

  • return 0

• node_count := number of nodes in the tree

• count_curr := 0

• current := root

• while current is not null, do

  • if current.left null, then

    • count_curr := count_curr + 1

    • if node_count mod 2 is not 0 and count_curr is same as(node_count + 1) /2, then

      • return previous.data

    • otherwise when node_count mod 2 is 0 and count_curr is same as(node_count/2) +1, then

      • return(previous.data + current.data) /2

    • previous := current

    • current := current.right

  • otherwise,

    • previous := current.left

    • while previous.right is not null and previous.right is not same as current, do

      • previous := previous.right

    • if previous.right is null, then

      • previous.right := current

      • current := current.left

    • otherwise,

      • previous.right := None

      • previous := previous

      • count_curr := count_curr + 1

      • if node_count mod 2 is not 0 and count_curr is same as(node_count + 1) / 2, then

        • return current.data

      • otherwise when node_count mod 2 is 0 and count_curr is same as(node_count / 2) + 1, then

        • return(previous.data+current.data) /2

      • previous := current

      • current := current.right

Example

Let us see the following implementation to get better understanding −

 Live Demo

class TreeNode:
   def __init__(self, data):
      self.data = data
      self.left = None
      self.right = None
def number_of_nodes(root):
   node_count = 0
   if (root == None):
      return node_count
   current = root
   while (current != None):
      if (current.left == None):
         node_count+=1
         current = current.right
      else:
         previous = current.left
         while (previous.right != None and previous.right != current):
            previous = previous.right
         if(previous.right == None):
            previous.right = current
            current = current.left
         else:
            previous.right = None
            node_count += 1
            current = current.right
   return node_count
def calculate_median(root):
   if (root == None):
      return 0
   node_count = number_of_nodes(root)
   count_curr = 0
   current = root
   while (current != None):
      if (current.left == None):
         count_curr += 1
         if (node_count % 2 != 0 and count_curr == (node_count + 1)//2):
            return previous.data
         elif (node_count % 2 == 0 and count_curr == (node_count//2)+1):
            return (previous.data + current.data)//2
         previous = current
         current = current.right
      else:
         previous = current.left
         while (previous.right != None and previous.right != current):
            previous = previous.right
         if (previous.right == None):
            previous.right = current
            current = current.left
         else:
            previous.right = None
            previous = previous
            count_curr+= 1
            if (node_count % 2 != 0 and count_curr == (node_count + 1) // 2 ):
               return current.data
            elif (node_count%2 == 0 and count_curr == (node_count // 2) + 1):
               return (previous.data+current.data)//2
            previous = current
            current = current.right
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(calculate_median(root))

Input

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)

Output

7