Average of Levels in Binary Tree

문제

Given the root of a binary tree, return the average value of the nodes on each level in the form of an array. Answers within 10-5 of the actual answer will be accepted.

풀이

• 같은 레벨에 존재하는 노드의 평균값을 구한다.
• 같은 레벨에 있는 값을 저장하는 map, 같은 레벨에 있는 노드의 개수를 저장하는 map 변수를 둔다.
• bfs로 탐색하면서 값들을 map에 저장한다.
• 탐색 완료후, 각 레벨의 평균을 구하여 리턴한다.

수도 코드

depths = key: level, value: sum(node.val)인 map
visited = key: level, value: count(node)인 map

q = (level, node)로 저장되는 큐
while q가 빌때까지:
	now = q.popleft()
    
    if depths에 now.level이 없다면:
    	depths에 저장
        visited에 저장
    else: # 존재
    	depths[level] 갱신
        visited[level] 갱신
    
    if now.node.left != null:
    	q에 (level + 1, now.node.left) 추가
    if now.node.right != null:
    	q에 (level + 1, now.node.right) 추가

answer = []
for level, total in depths.items():
	answer.append(total / visited[level])

Solution(Runtime: 11ms)

import java.util.*;


class Solution {
    static class Pair {
        public int level;
        public TreeNode node;

        public Pair(int l, TreeNode n) {
            level = l;
            node = n;
        }
    }

    public List<Double> averageOfLevels(TreeNode root) {
        Map<Integer, Double> depths = new HashMap<>();
        Map<Integer, Integer> visited = new HashMap<>();
        Queue<Pair> q = new LinkedList<>();
        q.add(new Pair(0, root));

        while (!q.isEmpty()) {
            Pair now = q.poll();

            if (depths.containsKey(now.level)) {
                double total = depths.get(now.level) + (double) now.node.val;
                depths.put(now.level, total);
                visited.put(now.level, visited.get(now.level) + 1);
            } else {
                depths.put(now.level, (double) now.node.val);
                if (!visited.containsKey(now.level)) {
                    visited.put(now.level, 0);
                }
                visited.put(now.level, visited.get(now.level) + 1);
            }

            if (now.node.left != null) {
                q.add(new Pair(now.level + 1, now.node.left));
            }

            if (now.node.right != null) {
                q.add(new Pair(now.level + 1, now.node.right));
            }
        }

        return depths.entrySet()
                        .stream()
                        .map(e -> {
                            double total = e.getValue();
                            return total / visited.get(e.getKey());
                        })
                        .collect(Collectors.toList());
    }
}

개선한 코드(Runtime: 2ms)

import java.util.*;


class Solution {

    public List<Double> averageOfLevels(TreeNode root) {
        List<Double> answer = new ArrayList<>();
        Queue<TreeNode> q = new LinkedList<>();
        q.add(root);

        while (!q.isEmpty()) {
            int qLen = q.size();
            double total = 0;
            
            for (int i = 0; i < qLen; i++) {
                TreeNode now = q.poll();
                total += now.val;

                if (now.left != null) {
                    q.add(now.left);
                }
                if (now.right != null) {
                    q.add(now.right);
                }
            }
            answer.add(total / qLen);
        }
        return answer;
    }
}

시간복잡도는 트리의 속한 노드 개수 n을 순회하므로 $O(n)$이 소요된다. 현재 시점에 큐에 존재하는 node는 같은 레벨이므로 map에 따로 저장할 필요 없이 현재 q의 크기만큼 순회한 후, answer에 평균을 구해서 저장하면 된다.