题目
来源
题目说明
给你一棵 完全二叉树 的根节点 root ,求出该树的节点个数。
完全二叉树
的定义如下:在完全二叉树中,除了最底层节点可能没填满外,其余每层节点数都达到最大值,并且最下面一层的节点都集中在该层最左边的若干位置。若最底层为第 h
层,则该层包含
1
∼
2
h
1 \sim 2h
1∼2h 个节点。
示例 1:
输入:root = [1,2,3,4,5,6]
输出:6
示例 2:
输入:root = []
输出:0
示例 3:
输入:root = [1]
输出:1
提示:
- 树 中 节 点 的 数 目 范 围 是 [ 0 , 5 ∗ 104 ] 树中节点的数目范围是[0, 5 * 104] 树中节点的数目范围是[0,5∗104]
- 0 ≤ N o d e . v a l ≤ 5 ∗ 104 0 \leq Node.val \leq 5 * 104 0≤Node.val≤5∗104
- 题目数据保证输入的树是 完全二叉树
进阶:遍历树来统计节点是一种时间复杂度为 O ( n ) O(n) O(n) 的简单解决方案。你可以设计一个更快的算法吗?
解题
方法:二分查找 + 位运算
对于任意二叉树,都可以通过广度优先搜索或深度优先搜索计算节点个数,时间复杂度和空间复杂度都是 O ( n ) O(n) O(n),其中 n n n是二叉树的节点个数。这道题规定了给出的是完全二叉树,因此可以利用完全二叉树的特性计算节点个数。
规定根节点位于第 0 0 0 层,完全二叉树的最大层数为 h h h。根据完全二叉树的特性可知,完全二叉树的最左边的节点一定位于最底层,因此从根节点出发,每次访问左子节点,直到遇到叶子节点,该叶子节点即为完全二叉树的最左边的节点,经过的路径长度即为最大层数 h h h。
当 0 ≤ i < h 0 \le i < h 0≤i<h 时,第 i i i 层包含 2 i 2^i 2i 个节点,最底层包含的节点数最少为 1 1 1,最多为 2 h 2^h 2h 。
当最底层包含
1
1
1 个节点时,完全二叉树的节点个数是
∑
i
=
0
h
−
1
2
i
+
1
=
2
h
\sum_{i=0}^{h-1}2^i + 1 = 2^h
i=0∑h−12i+1=2h
当最底层包含
2
h
2^h
2h 个节点时,完全二叉树的节点个数是
∑
i
=
0
h
2
i
=
2
h
−
1
\sum_{i=0}^{h}2^i = 2^h - 1
i=0∑h2i=2h−1
因此对于最大层数为 h h h的完全二叉树,节点个数一定在 [ 2 h , 2 h + 1 − 1 ] [2^h,2^{h+1}-1] [2h,2h+1−1]的范围内,可以在该范围内通过二分查找的方式得到完全二叉树的节点个数。
具体做法是,根据节点个数范围的上下界得到当前需要判断的节点个数 k k k,如果第 k k k个节点存在,则节点个数一定大于或等于 k k k,如果第 k k k个节点不存在,则节点个数一定小于 k k k,由此可以将查找的范围缩小一半,直到得到节点个数。
如何判断第
k
k
k 个节点是否存在呢?如果第
k
k
k个节点位于第
h
h
h 层,则
k
k
k 的二进制表示包含
h
+
1
h+1
h+1 位,其中最高位是
1
1
1,其余各位从高到低表示从根节点到第
k
k
k个节点的路径,
0
0
0表示移动到左子节点,
1
1
1 表示移动到右子节点。通过位运算得到第
k
k
k个节点对应的路径,判断该路径对应的节点是否存在,即可判断第
k
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 ) O(\log^2 n) O(log2n),其中 n n n是完全二叉树的节点数。
首先需要 O ( h ) O(h) O(h)的时间得到完全二叉树的最大层数,其中 hh 是完全二叉树的最大层数。
使用二分查找确定节点个数时,需要查找的次数为 O ( log 2 h ) = O ( h ) O(\log 2^h)=O(h) O(log2h)=O(h),每次查找需要遍历从根节点开始的一条长度为 h h h 的路径,需要 O ( h ) O(h) O(h) 的时间,因此二分查找的总时间复杂度是 O ( h 2 ) O(h^2) O(h2)。
因此总时间复杂度是 O ( h 2 ) O(h^2) O(h2)。由于完全二叉树满足 2 h ≤ n < 2 h + 1 2^h \le n < 2^{h+1} 2h≤n<2h+1 ,因此有 O ( h ) = O ( log n ) O(h)=O(\log n) O(h)=O(logn), O ( h 2 ) = O ( log 2 n ) O(h^2) = O(\log^2 n) O(h2)=O(log2n)。 -
空间复杂度: O ( 1 ) O(1) O(1)。只需要维护有限的额外空间。