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
