[LeetCode By Go 13]637. Average of Levels in Binary Tree

题目

Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.

Example 1:
Input:

  3
   / \
  9  20
    /  \
   15   7

Output: [3, 14.5, 11]
Explanation:
The average value of nodes on level 0 is 3, on level 1 is 14.5, and on level 2 is 11. Hence return [3, 14.5, 11].

Note:
The range of node's value is in the range of 32-bit signed integer.

分析

二叉树层序遍历

代码

averageOfLevels.go

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */


func averageOfLevel(nodeSlice []*TreeNode) (nextNodeSlice []*TreeNode, avg float64) {
    var sum float64

    num := len(nodeSlice)
    for _, node := range nodeSlice {
        sum += float64(node.Val)
        if nil != node.Left {
            nextNodeSlice = append(nextNodeSlice, node.Left)
        }
        if nil != node.Right {
            nextNodeSlice = append(nextNodeSlice, node.Right)
        }
    }

    avg = sum / float64(num)

    return nextNodeSlice, avg
}

func averageOfLevels(root *TreeNode) []float64 {
    var ret []float64
    var avg float64
    nextNodeSlice := []*TreeNode{root}

    for ; len(nextNodeSlice) > 0; {
        nextNodeSlice, avg = averageOfLevel(nextNodeSlice)
        ret = append(ret, avg)
    }

    return ret
}

测试

averageOfLevels_test.go

package _637_Average_Levels_Binary_Tree

import "testing"


func InitNode(val int, left *TreeNode, right *TreeNode) (ret *TreeNode){
    ret = new(TreeNode)
    ret.Val = val
    ret.Left = left
    ret.Right = right

    return ret
}

func SliceEqual(s1, s2 []float64) bool {
    len1 := len(s1)
    len2 := len(s2)

    if len1 != len2 {
        return false
    }
    for i := 0; i < len1; i++ {
        if s1[i] != s2[i] {
            return false
        }
    }

    return true
}
func TestAverageOfLevels(t *testing.T) {
    l3_1 := InitNode(15, nil, nil)
    l3_2 := InitNode(7, nil, nil)

    l2_1 := InitNode(9, nil, nil)
    l2_2 := InitNode(20, l3_1, l3_2)

    l1 := InitNode(3, l2_1, l2_2)

    ret := AverageOfLevels(l1)

    want := []float64{3, 14.5, 11}

    ok := SliceEqual(ret, want)

    if ok {
        t.Logf("pass")
    } else {
        t.Errorf("fail, want %+v, get %+v", want, ret)
    }
}