LeetCode222. 完全二叉树的节点个数Golang版

LeetCode222. 完全二叉树的节点个数Golang版

1. 问题描述

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

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

2. 思路

2.1. 递归思路

2.2. 位运算+二分查找

3. 代码

3.1. 递归代码

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func countNodes(root *TreeNode) int {
    if root == nil {
        return 0
    }
    
    left := root.Left
    right := root.Right
    leftHeight := 0
    rightHeight := 0

    for left != nil {
        leftHeight++
        left = left.Left
    }
    for right != nil {
        rightHeight++
        right = right.Right
    }
    if leftHeight == rightHeight {
        return 2 << leftHeight - 1
    }

    count := countNodes(root.Left) + countNodes(root.Right) + 1
    return count
}

3.2. 迭代代码

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

    node := root
    level := 0
    for node.Left != nil {
        level++
        node = node.Left
    }
    // fmt.Println(level)
    low := 1 << level
    high := 1 << (level + 1) - 1
    for low < high {
        // 必须+1，否则会话会出现low不更新，TLE(超时)
        mid := low + (high - low + 1) / 2
        if exists(root, level, mid) {
            low = mid
        } else {
            high = mid - 1
        }
    }
    return low
}

func exists(root *TreeNode, level int, k int) bool {
    bits := 1 << (level - 1)
    node := root

    for node != nil && bits > 0 {
        if bits & k == 0 {
            node = node.Left
        } else {
            node = node.Right
        }
        bits >>= 1
    }

    if node != nil {
        return true
    }
    return false
}