C++로 이진 트리 레벨별 평균값 구하기

문제 개요

비어 있지 않은(non-empty) 이진 트리가 하나 주어졌다고 가정해 보겠습니다. 우리가 해야 할 일은 트리의 각 레벨(level, 층)에 있는 노드 값들의 평균을 계산한 뒤, 그 평균값들을 배열 형태로 반환하는 것입니다.

예를 들어 입력으로 다음과 같은 이진 트리가 주어진다면,

출력은 [3, 14.5, 11]이 됩니다.

첫 번째 레벨에는 루트 노드 3만 존재하므로 평균은 3입니다.
두 번째 레벨의 노드 값은 9와 20이므로 평균은 (9 + 20) / 2 = 14.5입니다.
세 번째 레벨의 노드 값은 15와 7이므로 평균은 (15 + 7) / 2 = 11입니다.

해결 접근 방법

이 문제는 BFS(너비 우선 탐색), 즉 레벨 순서대로 트리를 탐색하는 방식으로 깔끔하게 해결할 수 있습니다. 큐(queue)를 활용하면 같은 레벨에 속한 노드들을 한꺼번에 처리할 수 있습니다. 구체적인 알고리즘 단계는 다음과 같습니다.

평균값을 저장할 배열 result를 정의합니다.
큐 q를 하나 정의하고 루트(root) 노드를 삽입합니다.
큐가 빌 때까지 다음 과정을 반복합니다.
- n := 현재 큐의 크기 (현재 레벨의 노드 수)
- 임시 배열 temp를 정의합니다.
- n이 0이 될 때까지 반복합니다.
  - t := 큐의 맨 앞 요소
  - t의 값을 temp에 추가
  - 큐에서 해당 요소를 제거
  - t의 왼쪽 자식이 null이 아니면 큐에 삽입
  - t의 오른쪽 자식이 null이 아니면 큐에 삽입
  - n을 1 감소
- temp의 크기가 1이면 result의 끝에 temp[0]을 그대로 추가합니다.
- temp의 크기가 1보다 크면 모든 값을 더한 뒤(sum := 0에서 시작해 temp의 모든 요소를 누적) 노드 수로 나누어 평균을 result에 추가합니다.
모든 레벨의 처리가 끝나면 result를 반환합니다.

이 알고리즘의 시간 복잡도와 공간 복잡도는 모두 O(n)입니다. 여기서 n은 트리의 전체 노드 수로, 각 노드를 정확히 한 번씩 방문하기 때문입니다.

예제 코드

아래는 위 알고리즘을 C++로 구현한 전체 코드입니다.

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class TreeNode{
    public:
        int val;
        TreeNode *left, *right;
        TreeNode(int data){
            val = data;
            left = NULL;
            right = NULL;
        }
};
void insert(TreeNode **root, int val){
    queue<TreeNode*> q;
    q.push(*root);
    while(q.size()){
        TreeNode *temp = q.front();
        q.pop();
        if(!temp->left){
            if(val != NULL)
                temp->left = new TreeNode(val);
            else
                temp->left = new TreeNode(0);
            return;
        }
        else{
            q.push(temp->left);
        }
        if(!temp->right){
            if(val != NULL)
                temp->right = new TreeNode(val);
            else
                temp->right = new TreeNode(0);
            return;
        }
        else{
            q.push(temp->right);
        }
    }
}
TreeNode *make_tree(vector<int> v){
    TreeNode *root = new TreeNode(v[0]);
    for(int i = 1; i<v.size(); i++){
        insert(&root, v[i]);
    }
    return root;
}
class Solution{
public:
    vector<float> averageOfLevels(TreeNode *root){
        vector<float> result;
        queue<TreeNode*> q;
        q.push(root);
        while (!q.empty()) {
            int n = q.size();
            vector<float> temp;
            while (n) {
                TreeNode* t = q.front();
                temp.push_back(t->val);
                q.pop();
                if (t->left && t->left->val != 0)
                    q.push(t->left);
                if (t->right && t->right->val != 0)
                    q.push(t->right);
                    n--;
            }
            if (temp.size() == 1)
                result.push_back(temp[0]);
            else if (temp.size() > 1) {
                double sum = 0;
                for (int i = 0; i < temp.size(); i++) {
                    sum += temp[i];
                }
                result.push_back(sum / temp.size());
            }
        }
        return result;
    }
};
main(){
    Solution ob;
    vector<int> v = {3,9,20,NULL,NULL,15,7};
    TreeNode *root = make_tree(v);
    print_vector(ob.averageOfLevels(root));
}

입력

{3,9,20,NULL,NULL,15,7}

출력

[3, 14.5, 11]