二叉树9——完全二叉树的节点个数

题目

来源

题目说明

给你一棵 完全二叉树 的根节点 root ，求出该树的节点个数。

完全二叉树 的定义如下：在完全二叉树中，除了最底层节点可能没填满外，其余每层节点数都达到最大值，并且最下面一层的节点都集中在该层最左边的若干位置。若最底层为第 h 层，则该层包含 $1\sim 2h$ 个节点。

示例 1：

输入：root = [1,2,3,4,5,6]
输出：6

示例 2：

输入：root = []
输出：0

示例 3：

输入：root = [1]
输出：1

提示：

• $树中节点的数目范围是[0,5\ast 104]$
• $0\le Node.val\le 5\ast 104$
• 题目数据保证输入的树是 完全二叉树

进阶：遍历树来统计节点是一种时间复杂度为 $O(n)$ 的简单解决方案。你可以设计一个更快的算法吗？

解题

方法：二分查找 + 位运算

对于任意二叉树，都可以通过广度优先搜索或深度优先搜索计算节点个数，时间复杂度和空间复杂度都是 $O(n)$，其中 $n$是二叉树的节点个数。这道题规定了给出的是完全二叉树，因此可以利用完全二叉树的特性计算节点个数。

规定根节点位于第 $0$ 层，完全二叉树的最大层数为 $h$。根据完全二叉树的特性可知，完全二叉树的最左边的节点一定位于最底层，因此从根节点出发，每次访问左子节点，直到遇到叶子节点，该叶子节点即为完全二叉树的最左边的节点，经过的路径长度即为最大层数 $h$。

当 $0\le i<h$ 时，第$i$ 层包含 $2^i$ 个节点，最底层包含的节点数最少为 $1$，最多为 $2^h$ 。

当最底层包含 $1$ 个节点时，完全二叉树的节点个数是
$$\sum _{i=0}^{h-1}{2}^i+1=2^h$$
当最底层包含 $2^h$ 个节点时，完全二叉树的节点个数是
$$\sum _{i=0}^h{2}^i=2^h-1$$

因此对于最大层数为 $h$的完全二叉树，节点个数一定在 $[2^h,2^{h+1}-1]$的范围内，可以在该范围内通过二分查找的方式得到完全二叉树的节点个数。

具体做法是，根据节点个数范围的上下界得到当前需要判断的节点个数 $k$，如果第 $k$个节点存在，则节点个数一定大于或等于 $k$，如果第 $k$个节点不存在，则节点个数一定小于 $k$，由此可以将查找的范围缩小一半，直到得到节点个数。

如何判断第 $k$ 个节点是否存在呢？如果第 $k$个节点位于第 $h$ 层，则 $k$ 的二进制表示包含 $h+1$ 位，其中最高位是 $1$，其余各位从高到低表示从根节点到第 $k$个节点的路径，$0$表示移动到左子节点，$1$ 表示移动到右子节点。通过位运算得到第 $k$个节点对应的路径，判断该路径对应的节点是否存在，即可判断第$k$ 个节点是否存在。

代码

class Solution {
public:
    int countNodes(TreeNode* root) {
        if (root == nullptr) {
            return 0;
        }
        int level = 0;
        TreeNode* node = root;
        while (node->left != nullptr) {
            level++;
            node = node->left;
        }
        int low = 1 << level, high = (1 << (level + 1)) - 1;
        while (low < high) {
            int mid = (high - low + 1) / 2 + low;
            if (exists(root, level, mid)) {
                low = mid;
            } else {
                high = mid - 1;
            }
        }
        return low;
    }

    bool exists(TreeNode* root, int level, int k) {
        int bits = 1 << (level - 1);
        TreeNode* node = root;
        while (node != nullptr && bits > 0) {
            if (!(bits & k)) {
                node = node->left;
            } else {
                node = node->right;
            }
            bits >>= 1;
        }
        return node != nullptr;
    }
};

复杂度分析

• 时间复杂度：$O({\log }^{2}n)$，其中 $n$是完全二叉树的节点数。
  首先需要 $O(h)$的时间得到完全二叉树的最大层数，其中 hh 是完全二叉树的最大层数。
  使用二分查找确定节点个数时，需要查找的次数为 $O(\log 2^h)=O(h)$，每次查找需要遍历从根节点开始的一条长度为 $h$ 的路径，需要 $O(h)$ 的时间，因此二分查找的总时间复杂度是 $O(h^2)$。
  因此总时间复杂度是 $O(h^2)$。由于完全二叉树满足 $2^h\le n<2^{h+1}$ ，因此有 $O(h)=O(\log n)$，$O(h^2)=O({\log }^{2}n)$。

• 空间复杂度：$O(1)$。只需要维护有限的额外空间。