Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 풀어보는 이진 트리 최대 레벨 합 문제

문제 개요

이진 트리의 루트 노드가 주어졌을 때, 루트의 레벨은 1이고 자식 노드의 레벨은 2이며, 그 아래로 내려갈수록 레벨이 하나씩 증가합니다. 이때 해당 레벨에 있는 모든 노드 값의 합이 가장 큰 레벨 X 중에서 가장 작은 값을 반환해야 합니다.

예를 들어 다음과 같은 이진 트리가 있다고 가정해 보겠습니다.

C++로 풀어보는 이진 트리 최대 레벨 합 문제

각 레벨의 합을 계산해 보면 레벨 1의 합은 1, 레벨 2의 합은 7 + 0 = 7, 레벨 3의 합은 7 + (-8) = -1입니다. 따라서 합이 가장 큰 레벨은 2이므로 출력 결과는 2가 됩니다.

접근 방법

이 문제는 너비 우선 탐색(BFS)을 활용하면 효율적으로 해결할 수 있습니다. 큐(queue)를 사용해 레벨별로 노드를 순회하면서 각 레벨의 합을 계산하고, 그중 최댓값을 가지는 레벨을 추적하는 방식입니다. 구체적인 단계는 다음과 같습니다.

  • level := 1, sum := 루트 노드 r의 값, ansLevel := level, ansSum := sum으로 초기화합니다.
  • 큐 q를 정의하고 루트 노드 r을 삽입합니다.
  • q가 빌 때까지 다음 과정을 반복합니다.
    • capacity := q의 현재 크기
    • level을 1 증가시키고 sum := 0으로 초기화
    • capacity가 0이 될 때까지 반복:
      • node := q의 맨 앞 노드를 꺼낸 후 삭제
      • 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;
      }
};
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:
   int maxLevelSum(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(){
   vector<int> v = {1,7,0,7,-8,NULL,NULL};
   TreeNode *root = make_tree(v);
   Solution ob;
   cout <<ob.maxLevelSum(root);
}

입력

[1,7,0,7,-8,null,null]

출력

2