[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.

完全二叉树的节点个数。

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

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

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/count-complete-tree-nodes
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

思路一 - 递归 + 分治

比较容易写出来的思路是可以用递归。我参考了这个帖子。

因为一个完美二叉树的节点个数是2^h - 1，2的h次方 - 1，h是树的高度。如果根节点的左子树和右子树的高度一样，则说明当前这棵树一定为满二叉树。如果左子树和右子树的高度不一样，则一定有其中一棵子树不是满二叉树，就需要递归到更小的子树去看到底哪个子树不是满二叉树。注意这个思路的时间复杂度是 O(logn * logn)。

复杂度

时间O(logn * logn)
空间O(logn) - O(h) - 树的高度

代码

Java实现

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
	public int countNodes(TreeNode root) {
		int left = leftDepth(root);
		int right = rightDepth(root);
		// 如果根节点的左子树深度等于右子树深度，则说明整棵树是满二叉树
		// 否则就要递归去找到底哪个子树不是满二叉树
		if (left == right) {
			// return (int) Math.pow(2, left) - 1;
			return (int) (1 << left) - 1;
		} else {
			return 1 + countNodes(root.left) + countNodes(root.right);
		}
	}

	private int leftDepth(TreeNode root) {
		int depth = 0;
		while (root != null) {
			root = root.left;
			depth++;
		}
		return depth;
	}

	private int rightDepth(TreeNode root) {
		int depth = 0;
		while (root != null) {
			root = root.right;
			depth++;
		}
		return depth;
	}
}
// O(logn * logn)

思路二 - 二分查找

注意这个题的 tag 是 binary search 二分法，二分法是这道题最优解但是不好想到，这里我也给出二分法的解法，参考了这个帖子。思路是首先拿到树的最大深度，这个不难，只要不停地往左孩子遍历就行。其次，如果是一个满二叉树（complete binary tree），节点数是2^depth - 1。按照例子一，如果是一个高度为3的树，节点个数就是2^3 - 1 = 7；同时最后一层的节点数应该是2^(depth - 1) = 2^2 = 4。

但是如果只是一个完全二叉树，最后一层的节点数很有可能是缺失的，此时就需要用二分法找最后一层到底有几个节点了。这里需要写一个 helper 函数帮助做二分法，因为每一层的节点数是跟当前的深度有关所以当前层的节点数的范围应该在0 ~ (2^depth - 1)之间。所以二分法找的时候，比如一个高度为 3 的树如果是满的，理论上最低一层应该有 8 个节点，则去判断是否有第 4 个节点，若有则往右子树走，若没有则往左子树走。最后 helper 函数返回的是最后一层处于位置i的节点是否存在。

最后着重解释一下exist函数。还是给定了上下界 left = 0, right = 2^depth - 1。这个函数是为了检查某一个节点是否存在的。检查的方法其实是类似 DFS，从深度为 0 的地方开始，在已经知道节点 index 的情况下，带入 tree 去检查，如果 index 大于当前层中间的 node，则去右子树看，反之则去左子树看。

复杂度

时间O(logn * logn)
空间O(1)

代码

Java实现

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public int countNodes(TreeNode root) {
        // corner case
        if (root == null) {
            return 0;
        }
        int d = computeDepth(root);
        if (d == 0) {
            return 1;
        }

        // normal case
        int left = 1;
        int right = (int) Math.pow(2, d) - 1;
        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (exists(root, d, mid)) {
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return (int) Math.pow(2, d) - 1 + left;
    }

    private int computeDepth(TreeNode root) {
        int depth = 0;
        while (root.left != null) {
            depth++;
            root = root.left;
        }
        return depth;
    }

    private boolean exists(TreeNode root, int depth, int index) {
        int left = 0;
        int right = (int) Math.pow(2, depth) - 1;
        int mid;
        for (int i = 0; i < depth; i++) {
            mid = left + (right - left) / 2;
            if (index <= mid) {
                root = root.left;
                right = mid;
            } else {
                root = root.right;
                left = mid;
            }
        }
        return root != null;
    }
}