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.

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.

解法1:

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

class Solution(object):
    def averageOfLevels(self, root):
        """
        :type root: TreeNode
        :rtype: List[float]
        """
        # BFS
        if not root:
            return []
        ans = []
        nodes = [root]
        while nodes:
            val = 0.0
            nodes2 = []
            for node in nodes:
                val += node.val
                if node.left:  nodes2.append(node.left)
                if node.right: nodes2.append(node.right)                    
            ans.append(val/len(nodes))
            nodes = nodes2
        return ans

精简版：

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

class Solution(object):
    def averageOfLevels(self, root):
        """
        :type root: TreeNode
        :rtype: List[float]
        """
        # BFS
        if not root:
            return []
        ans = []        
        nodes = [root]
        while nodes:            
            ans.append(sum(node.val for node in nodes)/(len(nodes)+0.0))
            nodes = [x for node in nodes for x in (node.left, node.right)  if x]                            
        return ans

解法2: DFS

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

class Solution(object):
    def averageOfLevels(self, root):
        """
        :type root: TreeNode
        :rtype: List[float]
        """
        if not root:
            return []
        ans = []
        self.helper(root, ans, 0)        
        return [a["sum"]/a["count"] for a in ans]

    def helper(self, root, ans, ith_level):
        if not root:
            return
        if ith_level == len(ans):
            ans.append({"sum": root.val, "count": 1.0})
        else:
            ans[ith_level]["sum"] += root.val        
            ans[ith_level]["count"] += 1
        self.helper(root.left, ans, ith_level+1)
        self.helper(root.right, ans, ith_level+1)

使用一个数组记录tree每个level的sum。数组的下标是level值，数组的元素是包含sum和count的map。

关于python 生成器：

>>> a=[1,2,3]
>>> def p(a):
... print a
... for i in a:
... print i
...
>>> p(x for x in a)
<generator object <genexpr> at 0x1074ab960>
1
2
3

可以看到sum（xx）的空间复杂度是1

leetcode 637. Average of Levels in Binary Tree

推荐阅读