"""Solution for 'Average levels iin binary tree problem.

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

Input:
    3
   / \
  9  20
    /  \
   15   7

Output: [3, 14.5, 11]
"""


class TreeNode(object):
    """Node of tree with value, left and right children."""
    def __init__(self, value, left=None, right=None):
        self.val = value
        self.left, self.right = left, right


def average_levels(root):

    levels = [root.val]

    def get_average_for_level(node):
        """"""
        next_level = 0
        if node.left:
            next_level += node.left.val
        if node.right:
            next_level += node.right.val

        nonlocal levels
        if next_level:
            levels.append(next_level / 2)

        if node.left:
            get_average_for_level(node.left)
        if node.right:
            get_average_for_level(node.right)

    get_average_for_level(root)

    return levels


root = TreeNode(3, TreeNode(9), TreeNode(20, TreeNode(15), TreeNode(7)))
result = average_levels(root)
print(result)