9. Count Complete Tree Nodes

Link to the problem

Description

Given a complete binary tree, count the number of nodes.

Example

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

Idea

Recursion. Compare the height of the left subtree and right subtree. If they are equal, then the left subtree is full, recurse on the right subtree. Otherwise, the right subtree is full, recurse on the left subtree.

Solution

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

18 / 18 test cases passed.
Runtime: 43 ms