C++로 이진 트리에서 노드 합이 최소가 되는 레벨 찾기

이진 트리가 하나 주어져 있다고 가정해 보겠습니다. 루트 노드의 레벨은 1이며, 그 자식 노드들은 레벨 2, 그다음 세대는 레벨 3처럼 순차적으로 증가합니다. 이때 우리가 구해야 할 것은 특정 레벨 X에 존재하는 모든 노드 값의 합이 최솟값이 되도록 하는 가장 작은 레벨 X입니다.

예를 들어 다음과 같은 트리가 있다고 합시다.

2번째 레벨의 노드 값 합은 4 + (-10) = -6으로, 다른 어떤 레벨보다도 작습니다. 따라서 출력 결과는 2가 됩니다.

문제 해결 접근 방법

이 문제는 BFS(너비 우선 탐색)을 활용하면 효율적으로 해결할 수 있습니다. BFS는 트리를 레벨 단위로 순회하기 때문에 각 레벨별 노드 값의 합을 손쉽게 계산할 수 있습니다. 알고리즘의 진행 과정은 다음과 같습니다.

level := 1, sum := 루트 노드 r의 값으로 초기화하고, ansLevel := level, ansSum := sum으로 설정합니다.
큐 q를 선언한 뒤 루트 노드 r을 삽입합니다.
큐 q가 빌 때까지 다음을 반복합니다.
- capacity := 현재 큐의 크기 (현재 레벨의 노드 수)
- level을 1 증가시키고, sum := 0으로 초기화합니다.
- capacity가 0이 될 때까지 반복합니다.
  - node := 큐의 맨 앞(front) 노드를 꺼내고, 큐에서 제거합니다.
  - node의 오른쪽 자식이 존재하면 sum에 해당 값을 더하고, 오른쪽 자식을 큐에 삽입합니다.
  - node의 왼쪽 자식이 존재하면 sum에 해당 값을 더하고, 왼쪽 자식을 큐에 삽입합니다.
  - capacity를 1 감소시킵니다.
- 현재까지의 최솟값(ansSum)보다 sum이 작으면, ansSum := sum, ansLevel := level로 갱신합니다.
모든 레벨을 순회한 후 ansLevel을 반환합니다.

C++ 구현 예제

아래 코드를 통해 실제 구현을 확인해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class TreeNode{
    public:
        int val;
        TreeNode *left, *right;
        TreeNode(int data){
            val = data;
            left = NULL;
            right = NULL;
        }
};
class Solution {
    public:
    int solve(TreeNode* r) {
        int level = 1, sum = r->val;
        int ansLevel = level, ansSum = sum;
        queue <TreeNode*> q;
        q.push(r);
        while(!q.empty()){
            int capacity = q.size();
            level++;
            sum = 0;
            while(capacity--){
                TreeNode* node = q.front();
                q.pop();
                if(node->right){
                    sum += node->right->val;
                    q.push(node->right);
                }
                if(node->left){
                    sum += node->left->val;
                    q.push(node->left);
                }
            }
            if(ansSum>sum){
                ansSum = sum;
                ansLevel = level;
            }
        }
        return ansLevel;
    }
};
main(){
    TreeNode *root = new TreeNode(5);
    root->left = new TreeNode(4);
    root->right = new TreeNode(-10);
    root->left->right = new TreeNode(-2);
    root->right->left = new TreeNode(-7);
    root->right->right = new TreeNode(15);
    Solution ob;
    cout <<ob.solve(root);
}

입력

TreeNode *root = new TreeNode(5);
root->left = new TreeNode(4);
root->right = new TreeNode(-10);
root->left->right = new TreeNode(-2);
root->right->left = new TreeNode(-7);
root->right->right = new TreeNode(15);

출력

동작 원리 살펴보기

위 예제 트리에서 각 레벨별 노드 값의 합을 계산해 보면 다음과 같습니다.

레벨 1: 5
레벨 2: 4 + (-10) = -6
레벨 3: (-2) + (-7) + 15 = 6

세 레벨 중 합이 가장 작은 레벨은 -6을 기록한 레벨 2이므로, 프로그램은 2를 출력합니다.

이 알고리즘은 트리의 모든 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(N)(N은 전체 노드 수)이며, 큐에는 최대 한 레벨의 노드만 저장되므로 공간 복잡도는 O(W)(W는 트리의 최대 너비)입니다. 음수 값을 가진 노드가 포함된 트리에서도 정확하게 동작한다는 점이 이 접근법의 장점입니다.