Problem Description

Given the root of a complete binary tree, your task is to count the total number of nodes in the tree. A complete binary tree is defined as a binary tree in which every level, except possibly the last, is completely filled, and all nodes are as far left as possible.

• root is a reference to the root node of the tree (or null if the tree is empty).
• You must return the total number of nodes in the tree as an integer.
• The tree is guaranteed to be complete, and each node has a unique value.

Constraints:

• The number of nodes in the tree is in the range [0, 5 * 10^4].
• The tree's node values are unique.

Thought Process

The most straightforward way to solve this problem is to traverse the entire tree and count each node. This brute-force approach would work for any binary tree. However, the problem specifically states that the tree is complete, which gives us an opportunity to optimize.

In a complete binary tree, all levels are fully filled except possibly the last, and all nodes in the last level are as far left as possible. This means that the structure is predictable, and we can use this property to avoid traversing every single node.

By comparing the height of the leftmost and rightmost paths, we can sometimes determine if a subtree is perfect (completely filled), allowing us to calculate the number of nodes in that subtree directly using the formula 2^h - 1 (where h is the height).

This insight leads us to a much more efficient solution than simply counting nodes one by one.

Solution Approach

Let's break down the optimized algorithm step by step:

1. Calculate the leftmost and rightmost heights:
   • Traverse from the root to the deepest left node to get the left height.
   • Traverse from the root to the deepest right node to get the right height.
2. Check if the subtree is perfect:
   • If the left and right heights are equal, the subtree is a perfect binary tree.
   • The number of nodes can be calculated directly as 2^h - 1.
3. If not perfect, recursively count nodes in the left and right subtrees:
   • Apply the same logic to the left and right children.
   • The total count is 1 + count(left) + count(right).
4. Base case:
   • If the node is null, return 0.

This approach leverages the properties of a complete binary tree to skip entire levels when possible, making it much faster than a naive traversal.

Example Walkthrough

Consider the following complete binary tree:

        1
       / \
      2   3
     / \  /
    4  5 6
  

• Step 1: The leftmost height (from 1 → 2 → 4) is 3.
• Step 2: The rightmost height (from 1 → 3 → 6) is 3.
• Step 3: Since heights are equal, the tree is perfect up to the last level, but the last level is not completely filled. So we recursively check the left and right subtrees.
• Step 4: For subtree rooted at 2:
  • Leftmost and rightmost heights are both 2, so this subtree is perfect: 2^2 - 1 = 3 nodes (2, 4, 5).
• Step 5: For subtree rooted at 3:
  • Leftmost height is 2 (3 → 6), rightmost height is 1 (3 only).
  • Since heights are not equal, recursively count left and right children.
  • Node 6: both heights are 1, so it's a single node.
• Step 6: Sum up: root (1) + left subtree (3) + right subtree (2) = 6 nodes.

Time and Space Complexity

Brute-Force Approach:

• Time Complexity: O(n) because every node is visited once.
• Space Complexity: O(h) due to recursion stack, where h is the height of the tree.

• Time Complexity: O((log n)^2). At each recursive call, we compute the height in O(log n) time, and there are at most O(log n) recursive levels.
• Space Complexity: O(log n) for the recursion stack.

The optimized approach is much faster for large trees because it can skip counting entire subtrees when they're perfect.

Summary

By leveraging the properties of complete binary trees, we avoid unnecessary traversal and make our solution highly efficient. The key insight is to recognize when a subtree is perfect and use the formula for the number of nodes directly, drastically reducing the number of recursive calls. This elegant strategy results in a solution that is both concise and performant, especially for large trees.

Code Implementation

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def countNodes(self, root: TreeNode) -> int:
        def left_height(node):
            h = 0
            while node:
                h += 1
                node = node.left
            return h

        def right_height(node):
            h = 0
            while node:
                h += 1
                node = node.right
            return h

        if not root:
            return 0

        lh = left_height(root)
        rh = right_height(root)

        if lh == rh:
            return (1 << lh) - 1
        else:
            return 1 + self.countNodes(root.left) + self.countNodes(root.right)
      

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};

class Solution {
public:
    int leftHeight(TreeNode* node) {
        int h = 0;
        while (node) {
            h++;
            node = node->left;
        }
        return h;
    }
    int rightHeight(TreeNode* node) {
        int h = 0;
        while (node) {
            h++;
            node = node->right;
        }
        return h;
    }
    int countNodes(TreeNode* root) {
        if (!root) return 0;
        int lh = leftHeight(root);
        int rh = rightHeight(root);
        if (lh == rh) {
            return (1 << lh) - 1;
        } else {
            return 1 + countNodes(root->left) + countNodes(root->right);
        }
    }
};
      

// Definition for a binary tree node.
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

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

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {number}
 */
var countNodes = function(root) {
    function leftHeight(node) {
        let h = 0;
        while (node) {
            h++;
            node = node.left;
        }
        return h;
    }
    function rightHeight(node) {
        let h = 0;
        while (node) {
            h++;
            node = node.right;
        }
        return h;
    }

    if (!root) return 0;
    let lh = leftHeight(root);
    let rh = rightHeight(root);
    if (lh === rh) {
        return (1 << lh) - 1;
    } else {
        return 1 + countNodes(root.left) + countNodes(root.right);
    }
};