LeetCode 222. Count Complete Tree Nodes（二分查找）

LeetCode 222. Count Complete Tree Nodes

Medium

Given a complete binary tree, count the number of nodes.

Note:

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 2h nodes inclusive at the last level h.

Example:

Input:

    1
   / \
  2   3
 / \  /
4  5 6

Output: 6

题意

求完全二叉树的节点个数

思路

用dfs可以以O(n)的时间复杂度求解，但是这种解法并没有利用到完全二叉树的性质。
考虑完全二叉树的性质，除了最后一行之外是满的，最后一行的节点尽量靠左排列。因此，只要求出最后一行的最右边的节点位置，就可以直接计算得到完全二叉树的节点个数。
考虑使用二分查找，用二进制表示树的路径，从高位到低位分别表示从根节点出发的每个分岔，0位表示向左，1位表示向右，例如000表示

  /
 /
/

而101表示

  \
  /
  \

于是乎，对完全二叉树最后一行的节点的二分查找就转化为对[0, (1<<(h-1)-1)]之间的整数的二分查找（其中h为完全二叉树的深度），如果以该整数为代表的路径存在则缩小查找范围至右半边，否则缩小查找范围至左半边。
给定一个整数，判断对应路径是否存在的时间复杂度为O(h)即O(logn)，其中n为完全二叉树节点的个数。完全二叉树最后一行节点个数的上界是n/2，所以二分查找至多进行O(logn)次，因此算法总的时间复杂度是O(log^2(n)).

代码

/**
 * 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) {
        int i = 0, lh = 0, rh = 0;
        TreeNode * ptr = root;
        while (ptr) {
            ptr = ptr->left;
            ++i;
        }
        lh = i;
        i = 0;
        ptr = root;
        while (ptr) {
            ptr = ptr->right;
            ++i;
        }
        rh = i;
        if (lh == rh) { // full tree with height {@code lh}
            return (1 << lh) - 1;
        }
        int l = 0, r = (1 << rh) - 1, mid = 0;  // bitwise 0 for left, 1 for right
        while (r - l > 1) { // binary search
            mid = (l + r) / 2;
            if (findPath(root, mid, rh)) {
                l = mid;
            } else {
                r = mid - 1;
            }
        }
        if (findPath(root, r, rh)) {
            return (1 << rh) + r;
        } else {
            return (1 << rh) + l;
        }
    }
private:
    /**
    * check if the path with height {@code h} exists from {@code root}
    * which is represented by {@code path}
    * bitwise 0 for left and 1 for right
    */
    bool findPath(TreeNode * root, int path, int h) {
        int bit = (1<<(h-1)), i = 0;
        for (i=0; i<h; ++i) {
            if (path & bit) {
                root = root->right;
            } else {
                root = root->left;
            }
            if (!root) {
                return false;
            }
            bit >>= 1;
        }
        return true;
    }
};