package leetcode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;

/**
 * https://leetcode.com/problems/average-of-levels-in-binary-tree/
 */
public class Problem637 {
    public static class TreeNode {
        long val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }
    }

    public List<Double> averageOfLevels(TreeNode root) {
        Map<Integer, ValueCount> map = new HashMap<>();
        Queue<LevelNode> queue = new LinkedList<>();
        queue.add(new LevelNode(1, root));
        while (!queue.isEmpty()) {
            LevelNode current = queue.remove();
            if (!map.containsKey(current.level)) {
                map.put(current.level, new ValueCount(current.node.val, 1));
            } else {
                ValueCount valueCount = map.get(current.level);
                valueCount.count++;
                valueCount.value += current.node.val;
            }
            if (current.node.left != null) {
                queue.add(new LevelNode(current.level + 1, current.node.left));
            }
            if (current.node.right != null) {
                queue.add(new LevelNode(current.level + 1, current.node.right));
            }
        }
        List<Double> result = new ArrayList<>();
        for (int i = 1; i <= map.size(); i++) {
            ValueCount valueCount = map.get(i);
            result.add((double) valueCount.value / (double) valueCount.count);
        }
        return result;
    }

    private static class ValueCount {
        private long value;
        private int count;

        public ValueCount(long value, int count) {
            this.value = value;
            this.count = count;
        }
    }

    private static class LevelNode {
        private final int level;
        private final TreeNode node;

        public LevelNode(int level, TreeNode node) {
            this.level = level;
            this.node = node;
        }
    }

    public static void main(String[] args) {
        Problem637 prob = new Problem637();

        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(9);
        root.right = new TreeNode(20);
        root.right.left = new TreeNode(15);
        root.right.right = new TreeNode(7);
        System.out.println(prob.averageOfLevels(root)); // [3, 14.5, 11]
    }
}

@@ -309,20 +309,19 @@
         
-long
+int
  val;%0A  
@@ -2061,382 +2061,7 @@
    %7D
-%0A%0A    public static void main(String%5B%5D args) %7B%0A        Problem637 prob = new Problem637();%0A%0A        TreeNode root = new TreeNode(3);%0A        root.left = new TreeNode(9);%0A        root.right = new TreeNode(20);%0A        root.right.left = new TreeNode(15);%0A        root.right.right = new TreeNode(7);%0A        System.out.println(prob.averageOfLevels(root)); // %5B3, 14.5, 11%5D%0A    %7D
 %0A%7D%0A

Update problem 637