每日水题 637. 二叉树的层平均值(python实现)

给定一个非空二叉树, 返回一个由每层节点平均值组成的数组。

示例 1：

输入：
    3
   / \
  9  20
    /  \
   15   7
输出：[3, 14.5, 11]
解释：
第 0 层的平均值是 3 ,  第1层是 14.5 , 第2层是 11 。因此返回 [3, 14.5, 11] 。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def averageOfLevels(self, root: TreeNode) -> List[float]:
        # 水题，直接DFS或者BFS即可，时间复杂度O(n)
        # 这里在下就使用一个队列实现之
        from queue import Queue
        q = Queue()
        self.num = [] # 记录某一层的个数
        self.cnt = [] # 记录某一层之和
        root.h = 0
        q.put(root)
        while not q.empty():
            temp = q.get() # 出队列
            if len(self.num) < temp.h + 1:
                self.num.append(1)
                self.cnt.append(temp.val)
            else:
                self.num[temp.h] = self.num[temp.h] + 1
                self.cnt[temp.h] = self.cnt[temp.h] + temp.val
            if temp.left:
                temp.left.h = temp.h + 1
                q.put(temp.left)
            if temp.right:
                temp.right.h = temp.h + 1
                q.put(temp.right)
        for i in range(0,len(self.num)):
            self.cnt[i] /= self.num[i]
        return self.cnt