LeetCode 637 - Average of Levels in Binary Tree

Average of Levels in Binary Tree - LeetCode

 

이진트리의 root가 주어졌을 때 각 레벨의 평균값을 담은 List를 반환하라.

예를 들어 [3, 9, 29, null, null, 15, 7]의 트리라면 레벨 1은 3.00000, 레벨 2는 (9 + 29) / 2 = 14.5000, 레벨 3은 (15 + 7) / 2 = 11.00000 이 된다.

 

레벨별로 어떤 노드들이 존재하는지를 파악해야 한다.

큐를 통해 해결할 수 있을 것 같다.

1. 루트 노드를 큐에 넣는다.
2. 큐에서 노드를 꺼낸다.
3. left, right를 보고 자식의 개수를 기록한다.
4. 자식들을 큐에 넣는다.
5. 반복을 통해 자식의 개수만큼 큐에서 꺼내면서 평균을 계산한다.
6. 자식들을 큐에 넣고 다음 반복에서는 이번에 기록된 자식의 개수만큼 처리를 한다.

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public List<Double> averageOfLevels(TreeNode root) {
        Queue<TreeNode> q = new LinkedList<>();
        q.add(root);
        List<Double> result = new ArrayList<>();

        double sum = 0.0;
        int cnt = 1;
        while (!q.isEmpty()) {
            int count = 0;
            for (int i = 0; i < cnt; i++) {
                TreeNode node = q.poll();
                sum += node.val;
                if (node.left != null) {
                    q.add(node.left);
                    count++;
                }
                if (node.right != null) {
                    q.add(node.right);
                    count++;
                }     
            }
            result.add(sum / cnt);
            sum = 0;
            cnt = count;
        }

        return result;
    }
}