【leetcode练习·二叉树】计算完全二叉树的节点数

本文参考labuladong算法笔记[拓展：如何计算完全二叉树的节点数 | labuladong 的算法笔记]

如果让你数一下一棵普通二叉树有多少个节点，这很简单，只要在二叉树的遍历框架上加一点代码就行了。

但是，力扣第第 222 题「完全二叉树的节点个数」给你一棵完全二叉树，让你计算它的节点个数，你会不会？算法的时间复杂度是多少？

这个算法的时间复杂度应该是 O(logN∗logN)O(logN∗logN)，如果你心中的算法没有达到这么高效，那么本文就是给你写的。

关于「完全二叉树」和「满二叉树」等名词的定义，可以参考基础知识章节的 二叉树基础。

一、思路分析

现在回归正题，如何求一棵完全二叉树的节点个数呢？

# 输入一棵完全二叉树，返回节点总数
def countNodes(root: TreeNode) -> int:

如果是一个普通二叉树，显然只要向下面这样遍历一边即可，时间复杂度 O(N)O(N)：

def countNodes(root: TreeNode) -> int:
    if root == None:
        return 0
    return 1 + countNodes(root.left) + countNodes(root.right)

那如果是一棵满二叉树，节点总数就和树的高度呈指数关系：

def countNodes(root: TreeNode) -> int:
    h = 0
    # 计算树的高度
    while root:
        root = root.left
        h += 1
    # 节点总数就是 2^h - 1
    return 2 ** h - 1

完全二叉树比普通二叉树特殊，但又没有满二叉树那么特殊，计算它的节点总数，可以说是普通二叉树和完全二叉树的结合版，先看代码：

class Solution:
    def countNodes(self, root: TreeNode) -> int:
        l = root
        r = root
        hl = 0
        hr = 0
        # 沿最左侧和最右侧分别计算高度
        while l is not None:
            l = l.left
            hl += 1
        while r is not None:
            r = r.right
            hr += 1
        # 如果左右侧计算的高度相同，则是一棵满二叉树
        if hl == hr:
            return pow(2, hl) - 1
        # 如果左右侧的高度不同，则按照普通二叉树的逻辑计算
        return 1 + self.countNodes(root.left) + self.countNodes(root.right)

结合刚才针对满二叉树和普通二叉树的算法，上面这段代码应该不难理解，就是一个结合版，但是其中降低时间复杂度的技巧是非常微妙的。

二、复杂度分析

开头说了，这个算法的时间复杂度是 O(log⁡N×log⁡N)O(logN×logN)，这是怎么算出来的呢？

直觉感觉好像最坏情况下是 O(N×log⁡N)O(N×logN) 吧，因为之前的 while 需要 log⁡NlogN 的时间，最后要 O(N)O(N) 的时间向左右子树递归：

return 1 + countNodes(root.left) + countNodes(root.right);

关键点在于，这两个递归只有一个会真的递归下去，另一个一定会触发 hl == hr 而立即返回，不会递归下去。

为什么呢？原因如下：

一棵完全二叉树的两棵子树，至少有一棵是满二叉树：

看图就明显了吧，由于完全二叉树的性质，其子树一定有一棵是满的，所以一定会触发 hl == hr，只消耗 O(log⁡N)O(logN) 的复杂度而不会继续递归。

综上，算法的递归深度就是树的高度 O(log⁡N)O(logN)，每次递归所花费的时间就是 while 循环，需要 O(log⁡N)O(logN)，所以总体的时间复杂度是 O(log⁡N×log⁡N)O(logN×logN)。

所以说，「完全二叉树」这个概念还是有它存在的原因的，不仅适用于数组实现二叉堆，而且连计算节点总数这种看起来简单的操作都有高效的算法实现。