222 Count Complete Tree Nodes

Given a complete binary tree, count the number of nodes.
Definition of a complete binary tree from Wikipedia:
In a complete binary tree every level, except possibly the last, is completely filled, 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$$.

Solution:
iterative-V1 (328 ms)

class Solution {
public:
    int left_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->left;
            i += 1;
        }
        return i;
    }
    int right_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->right;
            i += 1;
        }
        return i;
    }
    int countNodes(TreeNode* root) {
        if ( root == NULL ) return 0;
        int res = 0;
        while ( root != NULL ) {
            int ll = left_height(root->left), lr = right_height(root->left);
            int rl = left_height(root->right), rr = right_height(root->right);
            if ( ll == rr ) return res + pow(2, ll+1) - 1;
            if ( ll == rl ) {
                res += pow(2, ll);
                root = root->right;
                continue;
            }
            if ( ll == lr ) return res + pow(2, ll) + pow(2, rr) - 1;
            res += pow(2, rr);
            root = root->left;
        }
        return res;
    }
};

iterative-V2 (280 ms)

class Solution {
public:
    int left_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->left;
            i += 1;
        }
        return i;
    }
    int right_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->right;
            i += 1;
        }
        return i;
    }
    int countNodes(TreeNode* root) {
        if ( root == NULL ) return 0;
        int res = 0;
        int ll = 0, lr = 0, rl = 0, rr = 0;
        while ( root != NULL ) {
            if ( ll == 0 ) ll = left_height(root->left);
            if ( lr == 0 ) lr = right_height(root->left);
            if ( rl == 0 ) rl = left_height(root->right);
            if ( rr == 0 ) rr = right_height(root->right);
            if ( ll == rr ) return res + pow(2, ll+1) - 1;
            if ( ll == rl ) {
                res += pow(2, ll);
                root = root->right;
                ll = rl-1;
                rr = rr-1;
                rl = 0;
                lr = 0;
                continue;
            }
            if ( ll == lr ) return res + pow(2, ll) + pow(2, rr) - 1;
            res += pow(2, rr);
            root = root->left;
            ll = ll - 1;
            rr = lr - 1;
            lr = 0;
            rl = 0;
        }
        return res;
    }
};

iterative-V3 (80 ms)

class Solution {
    public: 
    int countNodes(TreeNode* root) { 
        int count = 0;
        int l_lh = 0, r_lh = 0;
        TreeNode *tnode;

        while (root != NULL ){
            if ( l_lh == 0 ) {
                tnode = root->left;
                while ( tnode != NULL ) {l_lh += 1;tnode = tnode->left;}
            }
            if ( r_lh == 0 ) {
                tnode = root->right;
                while ( tnode != NULL ) {r_lh += 1;tnode = tnode->left;}
            }
            if ( l_lh == r_lh ) {
                count += 1<<l_lh;
                root = root->right;
                l_lh = r_lh - 1;
                r_lh = 0;
            }
            else {
                count += 1<<r_lh;
                root = root->left;
                l_lh = l_lh -1;
                r_lh = 0;
            }
        }
        return count;
    }
};

Recursive

/**
 * 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 left_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->left;
            i += 1;
        }
        return i;
    }
    int right_height(TreeNode* root) {
        int i = 0;
        while ( root != NULL ) {
            root = root->right;
            i += 1;
        }
        return i;
    }
    int countNodes(TreeNode* root) {
        if ( root == NULL ) return 0;
        int ll = left_height(root->left), lr = right_height(root->left);
        int rl = left_height(root->right), rr = right_height(root->right);
        if ( ll == rr ) return pow(2, ll+1) - 1;
        if ( ll == rl ) return pow(2, ll) + countNodes(root->right);
        if ( ll == lr ) return pow(2, ll) + pow(2, rr) - 1;
        return pow(2, rr) + countNodes(root->left);
    }
};

Notes Record four heights: ll, lr, rl, rr.
The following relation always hold:

$$ll \leq lr \leq rl \leq rr$$ and $$rr - ll \leq 1$$.

Consider following four situation:

• n, n, n, n:
  number of nodes = $$2^{(n+1)}-1$$
• n, n, n, n-1:
  number of nodes = $$2^n$$ + countNodes(root->right)
• n, n, n-1, n-1:
  number of nodes = $$2^n + 2^{n-1}-1$$
• n, n-1, n-1, n-1:
  number of nodes = countNodes(root->left) + $$2^{n-1}$$

Actually we only have to consider two heights: ll and lr. Consider two following situation:

• $$ll = lr$$:
  Left sub-tree is full, number of nodes = $$2^{ll}$$ + countNodes(root->right);
• $$ll=lr+1$$: Right sub-tree is full, number of nodes = $$2^{lr}$$ + countNodes(root->left);