python leetcode 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.

求二叉树中每层的均值。

515. Find Largest Value in Each Tree Row

513. Find Bottom Left Tree Value

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:

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

# 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):
        """
        :type root: TreeNode
        :rtype: List[float]
        """
        ans = []
        node = [root]
        while node:
            ans.append(sum([i.val for i in node]) / len(node))
            curr = []
            for each in node:
                if each.left:
                    curr.append(each.left)
                if each.right:
                    curr.append(each.right)
                node = curr
        return ans