Problem #222: Problem #222

Intuition & Strategy Breakdown

Let's look at the problem. We are given a **Complete Binary Tree**. In a complete tree, every level is full, except possibly the last... where nodes are filled from left to right. Our goal? Count the total nodes in less than O of N time.

Python Solution

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        if not root:
            return 0
        
        def get_left_height(node):
            h = 0
            while node:
                h += 1
                node = node.left
            return h
        
        def get_right_height(node):
            h = 0
            while node:
                h += 1
                node = node.right
            return h
        
        lh = get_left_height(root)
        rh = get_right_height(root)
        
        if lh == rh:
            return (1 << lh) - 1
        
        return 1 + self.countNodes(root.left) + self.countNodes(root.right)

Java Solution

class Solution {
    public int countNodes(TreeNode root) {
        if (root == null) return 0;
        
        int lh = getLeftHeight(root);
        int rh = getRightHeight(root);
        
        if (lh == rh) {
            return (1 << lh) - 1;
        }
        
        return 1 + countNodes(root.left) + countNodes(root.right);
    }
    
    private int getLeftHeight(TreeNode node) {
        int h = 0;
        while (node != null) {
            h++;
            node = node.left;
        }
        return h;
    }
    
    private int getRightHeight(TreeNode node) {
        int h = 0;
        while (node != null) {
            h++;
            node = node.right;
        }
        return h;
    }
}