[LeetCode]222. Count Complete Tree Nodes 左右分区算法解析

题目描述

LeetCode原题链接：222. Count Complete Tree Nodes

Given the root of a complete binary tree, return the number of the nodes in the tree.

According to Wikipedia, every level, except possibly the last, is completely filled in a complete binary tree, and all nodes in the last level are as far left as possible. It can have between 1 and 2h nodes inclusive at the last level h.

Design an algorithm that runs in less than O(n) time complexity.

Example 1:

Input: root = [1,2,3,4,5,6]
Output: 6

Example 2:

Input: root = []
Output: 0

Example 3:

Input: root = [1]
Output: 1

Constraints:

• The number of nodes in the tree is in the range [0, 5 * 104].
• 0 <= Node.val <= 5 * 104
• The tree is guaranteed to be complete.

题干分析

根据 Complet Binary Tree 的性质：

1. 从根节点到倒数第二层是满二叉树；
2. 最后一层（第n层）节点可以不足 2ⁿ 个且这些节点从左到右依次填充。

我们可以发现，对于任意一棵完全二叉树上的任意一个节点，其左右孩子节点必然是满二叉树or完全二叉树。比如下面这棵满二叉树，对于节点A，它的左孩子就是一棵满二叉树，右孩子是一棵完全二叉树；而对于B节点，其左右孩子都是满二叉树；而对于c节点，其左孩子是一棵完全二叉树，右孩子则是满二叉树。

而我们知道，对于一棵高度是 h 的满二叉树，它共拥有 2^h - 1 个节点。那我们就可以根据这个公式来计算题目要求的完全二叉树节点总和：不断分割出满二叉树，根据公式求出每一个小满二叉树的节点树然后累加。

先来写一个求高度函数

我们细分的每个子树可能是满二叉树，也可能是完全二叉树，它的高度应该以最左面节点所在层级为准（Complet Binary Tree 性质二）。因此，我们只需要用一个指针来深度遍历到最左下方的节点就可以了。可以用递归来解，每一轮高度都加1，如果当前节点是NULL则返回-1。

int getHeight(TreeNode* root) {
    return root ? 1 + getHeight(root -> left) : -1;
}

这样计算出来的结果就是只有一层的二叉树对应高度是0，二层的二叉树对应高度是1，三层的二叉树对应高度是2......

分割原始二叉树

这个解法的关键是如何分割原始的二叉树 —— 根据前面的分析， 我们可以根据左右子树的高度差来判断哪边是满二叉树。具体步骤我们以上图的那棵树来讲解。设起始位置为根节点：

step1: 计算得到从根节点到最左下方节点的高度为5

step2: 计算根节点的右孩子的高度为2

step3: 计算高度差，可以看出heightL是包含根节点的高度，heightR不包含根节点，像上面这种情况，如果 heightR = heightL - 1，那么根节点的左孩子一定是一棵满二叉树（最后一层L节点前的位置都是排满的），此时可以直接计算出 根节点 + 左孩子 的节点和为 1 + (2³ - 1) = 2³，即 2^heightL。此时，我们已经完成了一次分割，然后将指针从A移到其右孩子，进入下一轮循环。

step4: 新一轮循环，从C节点开始，重复步骤1、2，计算高度差时发现相差为2，也就是第二种情况，此时左孩子是完全二叉树，右孩子是满二叉树：

我们将根节点（这里是C节点）和其右孩子分割出来，计算其节点和：1 + (2¹ - 1) = 2¹，即2^{heightL - 1}，然后和上一轮结果累加。指针移向左孩子。

step5: 新一轮循环，重复上面的步骤，继续累加分割的子树的节点和，最终当指针为NULL时表明分割原始二叉树完毕，跳出循环，得到最终节点总数。

上述过程中的heightL并不需要每一轮循环都调用getHeight函数。我们每次移动指针都是从根节点移到其左孩子或右孩子，相当于向下移动了一层。因此，只要在每轮循环末尾将上一轮计算的heightL减1就可以了；也就是说，只需要在最初计算一次heightL就足够了！这个过程可以用while循环来实现：

int countNodes(TreeNode* root) {
    int h = getHeight(root), count = 0;
    while(root) {
        // 情况一：左孩子是满二叉树（右孩子可能是完全二叉树，也可能是满二叉树）
        if(getHeight(root -> right) == h - 1) {
            count += 1 << h; // 根节点 + 左孩子
            root = root -> right;
        }
        // 情况二：左孩子是完全二叉树（右孩子是满二叉树，也可能不存在）
        else {
            count += 1 << (h - 1); // 根节点 + 右孩子
            root = root -> left;
        }
        h--;
    }
    return count;
}

完整代码示例（c++)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int countNodes(TreeNode* root) {
        int h = getHeight(root), count = 0;
        while(root) {
            if(getHeight(root -> right) == h - 1) {
                count += 1 << h;
                root = root -> right;
            }
            else {
                count += 1 << (h - 1);
                root = root -> left;
            }
            h--;
        }
        return count;
    }
    int getHeight(TreeNode* root) {
        return root ? 1 + getHeight(root -> left) : -1;
    }
};