问题描述
给出一个完全二叉树，求出该树的节点个数。
说明：
完全二叉树的定义如下：在完全二叉树中，除了最底层节点可能没填满外，其余每层节点数都达到最大值，并且最下面一层的节点都集中在该层最左边的若干位置。若最底层为第 h 层，则该层包含 1~ 2h 个节点。
示例:
输入:
1
/ \
2 3
/ \ /
4 5 6
输出: 6
解题思路
显然深搜广搜都可以做，
深搜
| class Solution:
def countNodes(self, root: TreeNode) -> int:
if root is None:
return 0
return 1 + self.countNodes(root.left) + self.countNodes(root.right)
|
广搜
1
2
3
4
5
6
7
8
9
10
11
12
13
14 | class Solution:
def countNodes(self, root: TreeNode) -> int:
if root is None:
return 0
q = deque([root])
ans = 0
while q:
ans += 1
node = q.popleft()
if node.left:
q.append(node.left)
if node.right:
q.append(node.right)
return ans
|
但是这样就没有利用到 完全二叉树 这个条件。
对于满二叉树而言，第 \(h\) 层（根结点为第 0 层，编号为 1）的结点编号范围为 \(2^h\sim 2^{h+1}-1\)。
对于完全二叉树而言，最左边的结点一定是在最后一层。
所以可以根据最后一层的结点编号范围进行二分搜索。
| 1
0/ \1
2 3
0/ \1 0/
4 5 6
|
根据结点编号的二进制表示可以确定它的路径，如在上图中编号为 \(\boxed{6}\) 的二进制表示为 \(1\boxed{10}\)，从二进制表示的次高位开始，为 \(1\) 表示在右子树，为 \(0\) 表示在左子树。然后判断按照该路径到达的结点是否存在即可。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34 | class Solution:
def countNodes(self, root: TreeNode) -> int:
if root is None:
return 0
def ok(idx: int, h: int) -> bool:
mask = 1 << (h - 1) # 从次高位开始
node = root
while mask:
if mask & idx: # 为 1 表示在右子树
node = node.right
else:
node = node.left
mask >>= 1
return bool(node)
h = 0
node = root.left
while node:
h += 1
node = node.left
lo, hi = 1 << h, (1 << (h + 1)) - 1
while lo < hi:
# 因为 lo=mid 可能造成死循环，
# 故在只有两个元素时，让 mid=hi
mid = (lo + hi + 1) // 2
if ok(mid, h):
lo = mid
else:
hi = mid - 1
return lo
|