Leetcode 222: Count Complete Tree Nodes

Do only the recursion will cause TLE. Thus, at each node, count the heights of its left subtree and right subtree, if left height equals to the right height, the number of nodes of this root can be calculated in the mathematical way.

/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
int countNodes(TreeNode* root) {
if(!root)
return 0;
int hl=0,hr=0;
TreeNode *l=root->left,*r=root->right;
while(l)
{
l=l->left;
hl++;
}
while(r)
{
r=r->right;
hr++;
}
if(hl==hr)
return (1<<hl+1)-1;//equals to pow(2,hl+1)-1,increase the speed
return 1+countNodes(root->left)+countNodes(root->right);
}
};