Leetcode Count Complete Tree Nodes

Leetcode Count Complete Tree Nodes，本算法主要可以使用二分搜索加速计算出最后一层最右边缺少的结点数，通过二分法，计算量变为O(h * h)，其中h为树高，相对递归等方法的O(2^(h - 1))大大减少了计算量。相关算法如下：

#include<iostream>
#include<vector>
#include<queue>

/*
 * Get the number of node in a completely tree with O(h * h) time
 */
using namespace std;
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 == NULL) {
            return 0;
        }
        // Get the height of the completely tree
        int depth = getTheDepth(root, -1);
        // Get missed number of node in the h depth of the tree
        int less = getLowerCount(root);
        // Calculate the number of the tree.
        int sum = (1 << depth) - 1 - less;
        return sum;
    }
    //Get the depth of the tree, direction: -1(the left subtree), 0(the left
    //subtree of the root's left subtree), 1(the right subtree).
    int getTheDepth(TreeNode* root, int direction) {
        int depth = 0;
        TreeNode* cur = root;
        if (direction == 0 && cur != NULL) {
            cur = cur->right;
            depth ++;
        }
        while (cur != NULL) {
            depth ++;
            if (direction <= 0) {
                cur = cur->left;
            } else {
                cur = cur->right;
            }
        }
        return depth;
    }
    // Count the missed nodes
    int getLowerCount(TreeNode* root) {
        TreeNode* cur = root;
        int left = getTheDepth(cur, -1);
        int right = getTheDepth(cur, 1);
        int mid = 0;
        int sum = 0;
        while (left > right) {
            mid = getTheDepth(cur, 0);
            if (left == mid) {
                cur = cur->right;
                left --;
                right --;
            } else {
                cur = cur->left;
                sum = sum + (1 << (left - 2));
                left --;
                right = getTheDepth(cur, 1);
            }
        }
        return sum;
    }
};
// Sample input: ./a.out argv1 argv2 ...
int main(int argc, char* argv[]) {
    Solution so;
    // Make the complete tree
    queue<TreeNode*> nodes;
    TreeNode* root = NULL;
    TreeNode* cur;
    for (int i = 1; i < argc; i ++) {
        if (root == NULL) {
            root = new TreeNode(atoi(argv[i]));
            nodes.push(root);
        } else {
            cur = nodes.front();
            nodes.pop();
            cur->left = new TreeNode(atoi(argv[i]));
            nodes.push(cur->left);
            i ++;
            if (i < argc) {
                cur->right = new TreeNode(atoi(argv[i]));
                nodes.push(cur->right);
            }
        }
    }
    // Get result
    int re = so.countNodes(root);
    cout<<"result: "<<re<<endl;
    return 0;
}