leetcode-222 完全二叉树的节点个数

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

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

示例 1：

输入：root = [1,2,3,4,5,6]
输出：6

示例 2：

输入：root = []
输出：0

示例 3：

输入：root = [1]
输出：1

提示：

• 树中节点的数目范围是[0, 5 * 104]
• 0 <= Node.val <= 5 * 104
• 题目数据保证输入的树是 完全二叉树

go实现：

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func countNodes(root *TreeNode) int {

	// 如果根节点为空，即节点数为0，直接返回
	if root == nil {
		return 0
	}

	// 获取左右子树高度
	leftHeight := getHeight(root.Left)
	rightHeight := getHeight(root.Right)

	// 因为是完全二叉树，如果左子树和右子树高度相同，表示左子树所有节点完整，故可直接推算，此时只需对右子树遍历即可
	// 节点数为2的幂次（即2^树高度-1），使用位移并显性幂运算（math.pow）性能更好，故采用位移运算
	if leftHeight == rightHeight {
		return (1 << leftHeight) + countNodes(root.Right)
	} else {
		// 当左子树和右子树高度不一致时，因为题设已说明最下层节点均在左边，说明右子树为完整节点，但比左子树少一层，故可直接推算（采用位移计算），此时只需遍历左子树即可
		return countNodes(root.Left) + (1 << rightHeight)
	}
}

// 辅助函数，用于计算树高度
func getHeight(node *TreeNode) int {

	// 定义高度计算器，初始值为0
	height := 0

	// 当前节点不为空时，计算器加1，同时将当前节点替换为左子节点，备下一次检查
	for node != nil {
		height++
		node = node.Left
	}
	// 返回最终计算结果
	return height
}