问题描述
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 2^h 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 * 10^4].
- 0 <= Node.val <= 5 * 10^4
- The tree is guaranteed to be complete.
问题分析
由于问题要求复杂度低于O(n)的算法，第一反应是能否以某种“二分查找”的方式求解，思考无果；参考评论区方法后，发现思路还是没理清，这道题只需要把握住完全二叉树的定义和树的高度之间的关系，利用递归即可求解，具体做法如下:
- 计算出root节点的树高 height(root)
- 如果 height(root.Right) == h-1，表示root的左右子树高度是一样的；也就是最后一个节点在右子树中; 而此时的左子树是完全二叉树，高度为 h-1, 节点数为 2^(h-1); 递归求解右树
- 如果 2 中的高度不等，表示最后一个节点在左子树中；右子树为完全二叉树，高度为 h-2, 节点数为 2^(h-2); 递归求解左树
求树高的复杂度为 O(logN)，每次递归过程中还需要再求一次树高，因此整个算法复杂度为 O(logN)^2
解法
// 1. 计算出root节点的树高 height(root)
// 2. 如果 height(root.Right) == h-1，表示root的左右子树高度是一样的；也就是最后一个节点在右子树中; 而此时的左子树是完全二叉树，高度为 h-1, 节点数为 2^(h-1); 递归求解右树
// 3. 如果 2 中的高度不等，表示最后一个节点在左子树中；右子树为完全二叉树，高度为 h-2, 节点数为 2^(h-2); 递归求解左树
func countNodes(root *TreeNode) int {
h := height(root)
if h < 0 {
return 0
}
if height(root.Right) == h-1 {
// 高度为h(0-index)时，完全二叉树的节点数为 2^(h+1) - 1;
// 这里左子树高度为(h-1), 加上root节点数量，即为 2^(h-1+1)-1+1 = 2^h
return 1<<h + countNodes(root.Right)
} else {
return 1<<(h-1) + countNodes(root.Left)
}
}
func height(root *TreeNode) int {
if root == nil {
return -1
}
return 1 + height(root.Left)
}