🖋
LeetCode 637. Average of Levels in Binary Tree
はじめに
LeetCode 「637. Average of Levels in Binary Tree」の問題をGoで解きました。
問題
問題文
Given the root of a binary tree, return the average value of the nodes on each level in the form of an array. Answers within 10-5 of the actual answer will be accepted.
和訳
二分木のrootが与えられた場合、各階層のノードの平均値を配列形式で返します。実際の回答値から10の5乗以内の回答が受け入れられます。
例
3
/ \
9 20
/ \
15 7
Input: root = [3,9,20,null,null,15,7]
Output: [3.00000,14.50000,11.00000]
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].
3
/ \
9 20
/ \
15 7
Input: root = [3,9,20,15,7]
Output: [3.00000,14.50000,11.00000]
制約
- The number of nodes in the tree is in the range [1, 104].
- -231 <= Node.val <= 231 - 1
解答
本問題はレベルごとにノードをグループ化して平均を出すという性質上、DFSよりBFSの方が自然かつ効率的に解けます。
コード
func averageOfLevels(root *TreeNode) []float64 {
queue := []*TreeNode{root}
var averages []float64
for len(queue) > 0 {
// queue には常に「現在のレベルのノード」が入っている
// そのノード数を先に取得しておくことで、
// 次のレベルのノードを後ろに追加しつつ、今のレベルの処理が終わることを保証できる
levelSize := len(queue)
sum := 0
for i := 0; i < levelSize; i++ {
node := queue[0]
queue = queue[1:]
sum += node.Val
if node.Left != nil {
queue = append(queue, node.Left)
}
if node.Right != nil {
queue = append(queue, node.Right)
}
}
averages = append(averages, float64(sum) / float64(levelSize))
}
return averages
}
計算量
- 時間的計算量：O(n)
- n はノード数（各ノードを1回ずつ訪れる）
- 空間的計算量：O(w)
- w は木の最大幅（Queue に最大で一度に保持するノード数）
参考
二分木の問題は、以下の記事でも紹介していますので、参考にしてください。
Discussion