222. Count Complete Tree Nodes

Leetcode Tree Binary Search

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

Example:

Input: 
    1
   / \
  2   3
 / \  /
4  5 6

Output: 6

分析

这道题目看起来是非常简单的。要求计算完全二叉树的节点个数。这不是很简单吗？遍历每一个节点。

 private int count;
    public int countNodes(TreeNode root) {
        count = 0;
        traversal(root);
        return count;
    }

    private void traversal(TreeNode root) {
        if (root == null) return;
        count++;
        traversal(root.left);
        traversal(root.right);
    }

然而提交结果是超时。所以肯定不能用$O(n)$的算法了, 最好是$O(\log(n))$的算法。

int height(TreeNode root) {
    return root == null ? -1 : 1 + height(root.left);
}

public int countNodes(TreeNode root) {
    int num = 0, h = height(root);
    while (root != null) {
        if (height(root.right) == h - 1) {
            num += 1 << h;
            root = root.right;
        } else {
            num += 1 << h-1;
            root = root.left;
        }
        h--;
    }
    return num;
}

树的高度height(root)可以连续访问左子节点得到: 单一节点的树高度为0; 如果整棵树是空的，则高度为-1。 首先确认右子树的高度是不是整棵树的高度减去1(height(root.right) == h - 1)。

如果是，则意味着左子树和右子树有相同的高度，左子树为完全二叉树，最后一个节点肯定在右子树。由于左子树的高度为$h-1$，且是完全二叉树，所以节点个数为$2^h - 1$。左子树节点数，加上根节点数1，然后迭代的加上右子树的节点数，即为整棵树的节点数。

如果不是，说明右子树的高度为$h-2$，且左子树不是完全二叉树，最后一个节点在左子树。右子树的节点数$2^{h-1}-1$，加上根节点数1，然后迭代的加上左子树的节点数，即为整棵树的节点数。

求树的高度的时间复杂度为$O(\log(n))$，一共有$O(\log(n))$次循环，总的时间复杂度为$O(\log(n)^2)$。