222. Count Complete Tree Nodes

Problem

Given a complete binary tree, count the number of nodes.

Note:

Definition of a complete binary tree from Wikipedia: In a complete binary tree every level, except possibly the last, is completely filled, 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.

Example:


Input: 
    1
   / \
  2   3
 / \  /
4  5 6

Output: 6

Summary

Output the number of nodes in a complete binary tree.

Solution Approach

For this problem, actually, we can traverse the tree once in level order, and then add up the number of nodes at each level.

Code


package leetcode

/**
 * 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
	}
	queue := []*TreeNode{}
	queue = append(queue, root)
	curNum, nextLevelNum, res := 1, 0, 1
	for len(queue) != 0 {
		if curNum > 0 {
			node := queue[0]
			if node.Left != nil {
				queue = append(queue, node.Left)
				nextLevelNum++
			}
			if node.Right != nil {
				queue = append(queue, node.Right)
				nextLevelNum++
			}
			curNum--
			queue = queue[1:]
		}
		if curNum == 0 {
			res += nextLevelNum
			curNum = nextLevelNum
			nextLevelNum = 0
		}
	}
	return res
}