LeetCode 222. Count Complete Tree Nodes

Intuition

Use the constraint that input is Complete Binary Tree. Every level is full except the last level filled from left to right.

At each step calculate

1. left_h: height of the left subtree
2. right_h: height of the right subtree

compare them to determine where the "missing nodes" are.

Case 1: left_h == right_h

Conclusion: Left subtree is a perfect binary tree and missing nodes are in somewhere in the right subtree

1. Add count of left subtree + root 2^left_h
2. Recurse on the right subtree to count the remaining nodes

Case 2: left_h > right_h

Conclusion: Right subtree is a perfect binary tree

1. Add count of right subtree + root 2^right_h
2. Recurse on the left subtree to count the remaining nodes

Optimal Solution

class Solution:
	def countNodes(self, root):
		if not root:
			return 0
			
		def getDepth(node):
			depth = 0
			while node:
				node = node.left
				depth += 1
			return depth
			
		left_h = getDepth(node.left)
		right_h = getDepth(node.right)
		
		
		if left_h == right_h:
			return 2 ** left_h + self.countNodes(root.right)
		else:
			return 2 ** right_h + self.countNodes(root.left)

Time Complexity

Time Complexity O(log^2n)

Space Complexity O(log n)