222. Count Complete Tree Nodes Easy

Given the root of a complete binary tree, return the number of the nodes in the tree.

According to Wikipedia, every level, except possibly the last, is completely filled in a complete binary tree, and all nodes in the last level are as far left as possible. It can have between 1 and 2^h nodes inclusive at the last level h.

Design an algorithm that runs in less than O(n) time complexity.

Example 1:
Input: root = [1,2,3,4,5,6]
Output: 6

Example 2:
Input: root = []
Output: 0

Example 3:
Input: root = [1]
Output: 1

Approach

Input: The root node of a complete binary tree root.

Output: Calculate the number of nodes in the tree with a time complexity of less than O(n).

This problem belongs to Binary Tree Traversal problems.

We can use the structural properties of a "complete binary tree" to reduce the time complexity from O(n) to O((log n)^2):

• If the left height of a tree (depth going all the way to the left) equals the right height (depth going all the way to the right), it means it is a perfect binary tree, and the number of nodes is directly 2^h - 1.
• Otherwise, it is not full, recursively count: 1 + nodes in left subtree + nodes in right subtree.

Because a complete binary tree guarantees that its subtrees are either perfect or still complete, every level can be pruned by the quick check of "whether it's a perfect tree".

Implementation

python

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        # Define function to calculate left subtree height
        def leftHeight(node):
            h = 0  # Initialize height to 0
            while node:  # Go all the way down left
                h += 1  # Increment height
                node = node.left
            return h  # Return left subtree height

        # Define function to calculate right subtree height
        def rightHeight(node):
            h = 0  # Initialize height to 0
            while node:  # Go all the way down right
                h += 1  # Increment height
                node = node.right
            return h  # Return right subtree height
        
        # Calculate left and right subtree heights of current tree
        lh = leftHeight(root)
        rh = rightHeight(root)

        # If left height equals right height, it's a perfect binary tree
        if lh == rh:
            return 2 ** lh - 1  # Perfect binary tree nodes formula: 2^height - 1
        
        # If not a perfect binary tree, recursively sum nodes of left and right subtrees plus root
        return 1 + self.countNodes(root.left) + self.countNodes(root.right)

javascript

/*
 * @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;
    }

    const lh = leftHeight(root);
    const rh = rightHeight(root);

    if (lh === rh) return 2 ** lh - 1;

    return 1 + countNodes(root.left) + countNodes(root.right);
};

Complexity Analysis

• Time Complexity: O((log n)^2)
• Space Complexity: O(log n)

Links

222. Count Complete Tree Nodes (English)

222. 完全二叉树的节点个数 (Chinese)