[ Top Interview 150 ] 637. Average of Levels in Binary Tree

문제

637. 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.

• Input : 이진 트리의 root 노드
• Output : 각 node의 계층 별 평균

알고리즘 전략

• 알고리즘 기본 전략 : 회귀를 사용

  • 트리의 특징은 어떤 노드를 잡더라도 그 노드로 시작하는 트리 구조이하는 것이다.
  • 즉, 회귀 전략을 이용한다면 쉽게 코드를 구현할 수 있다.

• 알고리즘 전략

  • 특정 계층의 노드의 개수와 노드의 합을 저장할 수 있는 Average 객체를 정의한다.
  • 왼쪽 노드의 각 계층별 정보를 담은 List<Average> 객체와 오른쪽 노드의 각 계층별 정보를 담은 List<Average> 객체를 가져온다.
  • 노드의 개수와 노드의 합을 토대로 두 정보를 합치고, 자신 노드의 정보를 맨 앞에 넣은 후에 반환한다.

문제 답안

• Time complexity : $O(nlogn)$

class Solution {
    public List<Double> averageOfLevels(TreeNode root) {
        List<Average> averages = averageOfLevelsByAverage(root);

        List<Double> answer = new ArrayList<>();
        for (int i = 0; i < averages.size(); i++) {
            answer.add(averages.get(i).getAverage());
        }
        return answer;
    }

    private List<Average> averageOfLevelsByAverage(TreeNode root) {
        if (root == null) {
            return new LinkedList<>();
        }
        
        List<Average> left = averageOfLevelsByAverage(root.left);
        List<Average> right = averageOfLevelsByAverage(root.right);
        
        List<Average> longList = left.size() >= right.size() ? left : right;
        List<Average> shortList = left.size() >= right.size() ? right : left;
        
        for (int i = 0; i < shortList.size(); i++) {
            longList.get(i).concat(shortList.get(i));
        }
        longList.add(0, new Average(root.val));
        return longList;
    }
	
    // Average 객체 정의
    class Average {
        private long count = 1;
        private long sum;
        
        public Average(int element) {
            sum = element;
        }
        
        public void concat(Average average) {
            if (average == null) {
                return;
            }
            
            this.count += average.count;
            this.sum += average.sum;
        }
        
        public double getAverage() {
            return (double) sum / count;
        }
    }
}

모범 답안 (Queue 사용)

• 알고리즘 전략

  • Queue를 생성하고 root를 넣음
  • Queue가 비어있을 때까지 반복
    • 현 시점에 Queue에 있는 노드들의 평균 값을 구함
    • 하위 노드가 있다면, Queue에 넣는다.

• 위 답안보다 더 좋은 점

  • 재귀 형식을 사용하지 않아 더 함수 Stack 관리가 용이함 따라서 시간과 메모리 사용량이 줄어듦

• 배운 점

  • Tree 구조가 재귀 형식의 형태를 가지고 있다 해서, 알고리즘을 무조건 재귀 형식으로 짤 필요는 없음
  • 다른 방법도 충분히 고려해 볼 것

class Solution {
    public List<Double> averageOfLevels(TreeNode root) {

        List<Double> averages = new ArrayList<>();
        if (root == null) {
            return averages;
        }
        
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(root);
        
        while (!queue.isEmpty()) {
            int count = queue.size();
            double sum = 0l;
            for (int i = 0; i < count; i++) {
                TreeNode curr = queue.poll();
                sum = sum + (double) curr.val;
                if (curr.left != null) {queue.add(curr.left);}
                if (curr.right != null) {queue.add(curr.right);}
            }
            averages.add(sum / count);
        }
        return averages;
    }
}