solved 'Count Complete Tree Nodes' on leetcode https://leetcode.com/problems/count-complete-tree-nodes/

count_complete_bt_path.cc

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
 
// 532 ms, too slow...
// Another way of binary search by maintaining a path vector.
class Solution {
public:
    int countNodes(TreeNode* root) {
        auto iter = root;
        int height = 0;
        while (iter) {
            ++height;
            iter = iter->left;
        }
        if (height <= 1) {
            return height;
        }

        int last_level_capacity = pow(2, height - 1);
        int num_last_level = 0;

        // Record the path to second last level:
        // 0 for left, 1 for right.
        vector<char> path(height - 2, 0);

        // Follow the path, try to find a second-to-last node which only
        // has one child.
        TreeNode* second_last_node;
        for (int i = 0; i < height - 2; ++i) {
            // Binary search.
            path[i] = 1;
            last_level_capacity /= 2;

            second_last_node = followPath(root, path);
            if (!second_last_node->left && !second_last_node->right) {
                // No children, should go left.
                path[i] = 0;
            } else if (second_last_node->left && second_last_node->right) {
                // The node is full, continue going right.
                num_last_level += last_level_capacity;
                continue;
            } else {
                // Has only one child.
                num_last_level += 1 + last_level_capacity;
                // Directly calculate the size.
                goto RES;
            }
        }
        second_last_node = followPath(root, path);
        if (second_last_node->left) {
            ++num_last_level;
        }
        if (second_last_node->right) {
            ++num_last_level;
        }
        RES:
        return pow(2, height - 1) - 1 + num_last_level;
    }
private:
    static TreeNode* followPath(TreeNode* iter, const vector<char>& path) {
        for (char direction : path) {
            if (direction) {
                iter = iter->right;
            } else {
                iter = iter->left;
            }
        }
        return iter;
    }
};

count_complete_bt_topdown.cc

#include <cassert>
#include <iostream>
#include <queue>

using namespace std;

// Definition for a binary tree node.
struct TreeNode {
  int val;
  TreeNode *left;
  TreeNode *right;
  TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

// Inspired from https://leetcode.com/discuss/42785/72-ms-c-solution
// 76 ms, O(lgn lgn).
class Solution {
public:
  int countNodes(TreeNode* root) {
    int height = 0;
    for (auto iter = root; iter; iter = iter->left) {
      ++height;
    }
    if (height == 0) {
      return 0;
    }

    auto iter = root;
    int node_cnt = 1;
    for (int depth = 1; depth < height; ++depth) {
      // Go pass a level, node number doubled.
      node_cnt *= 2;
      // Calculate heigh of right sub tree.
      int right_height = 0;
      for (auto iter_r = iter->right; iter_r; iter_r = iter_r->left) {
        ++right_height;
      }

      if (right_height + depth == height) {
        // Right sub tree reaches the max height, meaning the 
        // left sub tree is full.
        iter = iter->right;
        ++node_cnt;
      } else {
        // Otherwise continue going left to find the node with sub trees
        // having equal heights.
        iter = iter->left;
      }
    }
    return node_cnt;
  }
};

// Test utility function to build a complete binary
// tree with specified size by level order traversal.
TreeNode* buildCompleteBinaryTree(int num) {
  if (num == 0) {
    return nullptr;
  }

  queue<TreeNode*> node_queue;
  TreeNode *root = new TreeNode(42);
  node_queue.push(root);
  --num;

  // Loop while `num` goes to 0 (joke :P
  while (num --> 0) {
    auto node = node_queue.front();
    if (!node->left) {
      node->left = new TreeNode(42);
      node_queue.push(node->left);
    } else {
      node->right = new TreeNode(42);
      node_queue.push(node->right);
      node_queue.pop();
    }
  }

  return root;
}

// Free memory space of binary tree in a recursive way.
void deleteBinaryTree(TreeNode *subroot) {
  if (!subroot) {
    return;
  }
  deleteBinaryTree(subroot->left);
  deleteBinaryTree(subroot->right);
  delete subroot;
}

int main() {
  Solution sol;
  TreeNode *root;
  int res_count;

  // Create some test cases, in a pseudo-random way.
  int inc = 1;
  for (int i = 0; i < 200; i += (inc++)) {
    root = buildCompleteBinaryTree(i);
    res_count = sol.countNodes(root);
    assert(res_count == i);
    deleteBinaryTree(root);
  }
}