[Leetcode] 637. Average of Levels in Binary Tree

오늘은 Binary Tree 의 레벨별 평균을 구하는 easy 단계의 문제를 풀어보겠다.

선정의 다른 이유는 없고 March Leetcode Challenge 2021 을 진행하는데에 3월 5일 (PST)에 출시된 문제이기 때문이다.

 

문제는 아래 링크를 참고하시면 된다.

leetcode.com/problems/average-of-levels-in-binary-tree/

Average of Levels in Binary Tree - LeetCode

 

루트노드부터 시작해서 레벨1, 루트노드의 첫번째 2개의 자식들이 레벨2, 그 다음 4개의 자식들이 레벨3, ... 등으로 뻗어나갈때, 각 레벨별로 평균을 구하라는 문제이다.

트리 문제이므로 BFS (Breadth First Search, 넓이 우선 탐색) 로든, DFS (Depth First Search, 깊이 우선 탐색) 로든 모두 풀 수 있다.

 

BFS와 DFS는 잘 아시리라 믿는다. 만약에 두 방식의 구현이 헷갈리신다면, 그것부터 완벽하게 숙지하고 문제풀이를 해야한다. 각 방식별로 Recursive (재귀) 방식과 Iterative (반복) 방식으로 구현할 수 있어야 한다. 특히 DFS는 Pre-Order (전위), In-Order (중위), Post-Order (후위) 3개의 방식 모두 익숙해져야 한다.

 

필자는 BFS로 풀었다. (DFS가 공간복잡도가 낮다)

 

전체코드를 복사&붙여넣기 하는 행위는 실력에 전혀 도움되지 않음을 명심하시길 바란다.

class Solution {
    public List<Double> averageOfLevels(TreeNode root) {
        List<Double> result = new ArrayList<>(); // ArrayList이든 LinkedList이든 이 문제는 크게 상관없음
        if (root == null) { // 엣지케이스
            return result; // null을 리턴하면 오류난다
        }
        LinkedList<TreeNode> queue = new LinkedList<>(); // BFS는 큐를 사용
        queue.add(root); // 루트를 큐에 삽입
        
        while (!queue.isEmpty()) { // 큐가 빌때까지, 루트를 넣었으니 아래에서 자식노드들을 넣어주면 됨
            int size = queue.size(); // 현재 큐의 사이즈(자식 노드)를 측정해서 반복하는 것이 핵심
            double sumOfLevel = 0; // 레벨별 합계를 저장하기 위한 변수 선언, 반복이 끝나면 초기화 됨
            for (int i = 0; i < size; i++) { // 자식노드의 갯수만큼 반복
                TreeNode node = queue.poll(); // 큐에서 하나의 노드를 꺼냄
                sumOfLevel += node.val; // 꺼낸 노드의 값을 누적해줌
                if (node.left != null) { // 꺼낸 노드의 왼쪽 자식이 null이 아니면
                    queue.add(node.left); // 큐에 넣어줌
                }
                if (node.right != null) { // 꺼낸 노드의 오른쪽 자식이 null이 아니면
                    queue.add(node.right); // 큐에 넣어줌
                }
            } // 반복이 끝나면
            result.add(sumOfLevel / size); // 레벨별로 누적된 값을 사이즈로 나눠주면 레벨별 평균
        }
        
        return result; // 결과 리턴
    }
}

 

Leetcode는 친절한 문제풀이 사이트이다. 설계가 필요함에도 알고리즘에 집중할 수 있게 일부 설계를 도와주기 때문이다.

(접힌글 참고) 

/**
 * 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) {
        
    }
}

 

하지만 인터뷰에서 이 질문을 받았다고 가정하고, 설계부터 살펴보기로 한다.

트리가 문제로 주어졌을때 가장 중요한 것은 이 트리가 그냥 트리인지, 이진트리인지, 이진탐색트리인지 물어봐야 한다.
혹시 문제에 알려줬다고 하더라도, 본인이 이해한 이진트리가 맞는지 인터뷰어에게 확인하는 과정도 좋은 행위이다.
ex. 이진트리라고 하셨는데, 이진트리는 각 노드가 최대 2개의 자식을 가지는 트리구조를 말씀하시는 것 맞죠?

당연하겠지만 각각의 차이를 이해하는 것이 자료구조 이해의 출발이니 인지하고 있기를 바란다.

만약에 이 글을 읽으시는 분이 Leetcode로 인터뷰/코딩테스트 연습을 시작했다고 하면 트리 문제의 클래스/구조체명은 TreeNode로 하시는 게 좋다.

이유는? Leetcode에선 모든 트리 문제의 클래스명이 TreeNode이니까...ㅎㅎㅎ 직관적이기도 하고요!

public class TreeNode {
    int val; // 정수값 저장
    TreeNode left; // 왼쪽 자식
    TreeNode right; // 오른쪽 자식
    TreeNode(int val) { // 생성자
    	this.val = val; 
    }
TreeNode(int val, TreeNode left, TreeNode right) { // 테스트할때 사용
    	this.val = val;
        this.left = left;
        this.right = right;
    }
}

TreeNode안에 정수값을 저장할 변수가 필요할테니 val를 선언한다.

그리고 이진트리이니까 좌우로 자식노드들을 가지는 객체를 선언해준다. (left, right)

Leetcode에서 생성자를 여러가지 뒀는데, 보통 정수값을 매개변수로 하는 생성자만 쓰다보니 다른 것은 제거했다.

나중에 테스트를 위해 left, right 자식을 갖는 생성자를 추가해줬다.

---

이제 레벨별 평균을 찾아가는 메서드를 추가할 예정이다.

Leetcode에서는 아래와 같은 메서드 시그니처와 리턴을 사용했는데, 이런 것들이 주어진다면 그대로 사용해야한다.

    public List<Double> averageOfLevels(TreeNode root)

일부 회사의 (구글) 경우에는 어떤 파라미터를 원하는지, 리턴은 어떻게 할 것인지에 대한 질문을 한다고 들었다.
물론 이유까지 대답을 해야한다고 한다. (ex. Map 또는 List를 원하면, 왜 그렇게 선택했는지, 어떤 장/단점이 있는지)

그런 것들을 생각하면서 문제해결에 접근하면 더욱 좋을 것 같다. 저한테 하는 말입니다.ㅋㅋ

---

이제 대충 틀을 갖추었으니 실제 코딩에 들어가보겠다.

리턴을 설정했으니 리턴해줄 값을 첫 부분에 선언해준다. Leetcode에서 List<Double>로 리턴하라고 했으니, 우리는 ArrayList 또는 LinkedList를 선언해서 사용하면 된다.

리턴 결과가 매우 클때는 LinkedList를 사용하는 것이 맞겠지만, 이진트리에서 레벨별로 평균을 저장하므로 실제 저장되는 갯수는 logN 이다. 그래서 딱히 상관은 없으니 편한 구현체를 선택하면 될 것 같다.

        List<Double> result = new ArrayList<>(); // ArrayList이든 LinkedList이든 이 문제는 크게 상관없음
        if (root == null) { // 엣지케이스
            return result; // null을 리턴하면 오류난다
        }

 

그리고 엣지케이스로 입력된 파라미터가 null 인지를 체크해준다.

입력 파라미터가 클래스인 경우에는 null 체크정도로 충분하다.

 

---

노드를 큐에 저장하기 위해서 큐를 선언한다.

순서대로 삽입할 것이므로 LinkedList가 어울린다.

문제에서 입력 노드의 갯수는 안나와있지만, 입력은 많다고 가정해야 속이 편하다. 

        LinkedList<TreeNode> queue = new LinkedList<>(); // BFS는 큐를 사용
        queue.add(root); // 루트를 큐에 삽입

선언부를 인터페이스로 사용해도 되고, 구현체로 사용해도 되지만, 해당 인터페이스의 메서드들만 사용할 수 있게 되므로 필자의 경우에는 구현체로 선언을 한다.

그리고 큐를 선언했으니 루트노드를 삽입 (add) 해준다. 

필자의 경우에는 offer 대신에 add를 즐겨쓴다. 차이는 큐가 꽉 찼을때 add는 Exception을 던져준다는 것이다.

 

이제 다음 레벨의 자식 노드가 있다면 큐에 넣어주고 계속 반복될 수 있게 해주는 것이다.

        while (!queue.isEmpty()) { // 큐가 빌때까지, 루트를 넣었으니 아래에서 자식노드들을 넣어주면 됨
            int size = queue.size(); // 현재 큐의 사이즈(자식 노드)를 측정해서 반복하는 것이 핵심
            double sumOfLevel = 0; // 레벨별 합계를 저장하기 위한 변수 선언, 반복이 끝나면 초기화 됨
            for (int i = 0; i < size; i++) { // 자식노드의 갯수만큼 반복
                TreeNode node = queue.poll(); // 큐에서 하나의 노드를 꺼냄
                sumOfLevel += node.val; // 꺼낸 노드의 값을 누적해줌
                if (node.left != null) { // 꺼낸 노드의 왼쪽 자식이 null이 아니면
                    queue.add(node.left); // 큐에 넣어줌
                }
                if (node.right != null) { // 꺼낸 노드의 오른쪽 자식이 null이 아니면
                    queue.add(node.right); // 큐에 넣어줌
                }
            } // 반복이 끝나면
            result.add(sumOfLevel / size); // 레벨별로 누적된 값을 사이즈로 나눠주면 레벨별 평균
        }

우선 현재 큐의 사이즈를 체크한 뒤에, 각 레벨별로 합계를 누적할 변수를 선언한다.

레벨별로 합계를 누적해서 사이즈 (자식노드의 수) 만큼 나누게 되면 평균이 나오는데, 이 평균을 result에 넣어주면 된다.

어차피 result가 List<Double> 형이고, 나눗셈 연산이 있어서 double은 숙명...

 

사이즈별로 반복을 시작하면서, 큐에 제일 먼저 들어온 노드부터 꺼내서 값을 누적해준다. 그리고 해당 노드에 자식들이 있다면 큐에 삽입해준다.

선입선출의 큐이기 때문에 가장 먼저 들어온 것부터 poll 시키고, 자식들은 들어오는 순서대로 차례를 기다린다.

 

[1,2,3] 이라는 트리를 기준으로 말로 디버깅을 해보면, 처음 루트노드 [1] 를 큐에 넣었으니까 첫번째 큐 사이즈는 1이 될 것이다. 1번의 반복을 통해서 루트노드의 자식이 있다면 7~12줄 코드에서 자식들 [2,3] 을 큐에 삽입할 수 있게 될 것이다. 레벨1의 반복이 끝났으니 누적된 값의 총계는 1이고 사이즈도 1이었으니 평균의 값은 1.00000 이 삽입된다.

반복문의 처음으로 돌아가면서 큐의 사이즈는 2 로 결정되고, 레벨별합계 변수는 0으로 초기화 된다.

사이즈별 반복에 의해 큐에서 poll을 하면 선입선출에 의해 [2] 노드가 선택되고, 해당 값을 레벨별합계에 누적한다. 자식들이 있으면 큐에 삽입을 하겠지만 더 이상 자식은 없다. 다시 사이즈별 반복으로 돌아가서 poll을 하고, 남은 [3] 노드가 선택된다. 3을 레벨별합계에 누적하고, 자식은 없으므로 사이즈별 반복은 종료된다.

레벨2의 반복의 결과로 레벨별합계는 5가 되고, 사이즈가 2이니 평균은 2.50000 가 된다.

큐에 노드가 더 이상 없으므로, 반복은 종료되고 [1.0000, 2.50000] 을 리턴하며 메서드콜은 종료된다.

---

만약 인터뷰 상황이라면 구현을 끝내고 테스트를 해봐야하는데, 기본적인 테스트케이스를 할 수 있을 정도의 시간까지 고려해야 한다.

    public static void main(String[] args) {
        // [1,2,3] 트리 구성
        //   1
        //  / \
        // 2   3
        List<Double> result = new ArrayList<>(Arrays.asList(1.00000, 2.50000));
        TreeNode test = new TreeNode(1, new TreeNode(2), new TreeNode(3));        
        assert averageOfLevels(test).equals(result);
    }

---

이 문제의 시간복잡도는 입력된 root의 전체를 한번 탐색하므로, 길이가 N 이라고 했을때, $O(N)$이 된다.

공간복잡도는 탐색 과정 속에 노드를 저장하기 위한 큐를 하나 선언했고, 결과를 담기 위한 리스트를 하나 선언했다.

큐의 경우에는 레벨별 노드의 갯수가 담긴다. 모든 자식들이 존재한다는 가정하에 1레벨 (root) 의 경우에는 큐에 1개가 담길 것이고, 반복문 안에서 비워졌다가 자식을 추가하는 부분에서 2개가 담길 것이다. 마찬가지로 하나씩 빠지면서 자식이 2개가 추가될 것이다.

그 결과는 레벨별 노드의 수로 수렴하게 된다. 레벨별 노드의 수가 M 이라고 하면 $O(M)$ 이 된다.

그리고 선언한 리스트의 경우에는 상술했듯이 logN 만큼 담긴다.

그러므로 이 문제의 공간복잡도는 $O(M)$이 될 것이다. logN은 M에 비해 작은 수라서 표기법에 의해 소거된다.

---

마지막으로 Follow-up Question으로 나올만한 것을 생각해보면

• 재귀적인 BFS로 구현이 가능한가?
• BFS 대신 DFS로 구현해보라

혹시 추가할만한 질문들이 있다면 제보바랍니다.

---