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

C++로 풀어보는 이진 트리 코인 분배 문제

이진 트리의 루트 노드가 주어졌을 때, 트리에는 총 N개의 노드와 N개의 코인이 존재합니다. 각 노드는 node.val만큼의 코인을 가지고 있으며, 모든 노드에 정확히 하나씩의 코인이 배분되도록 만들어야 합니다.

한 번의 이동(move)으로 인접한 두 노드 사이에서 단 하나의 코인을 옮길 수 있습니다. 이동 방향은 부모 노드에서 자식 노드로, 또는 자식 노드에서 부모 노드로 어느 쪽이든 가능합니다. 우리가 구해야 하는 것은 모든 노드에 코인이 하나씩 배치되도록 하는 데 필요한 최소 이동 횟수입니다.

문제 예시

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

  • 루트: 코인 0개
  • 왼쪽 자식: 코인 3개
  • 오른쪽 자식: 코인 0개

이 경우 정답은 3입니다. 왼쪽 자식에서 루트로 코인 2개를 보내는 데 각각 한 번씩, 총 2번의 이동이 필요하고, 이후 루트에서 오른쪽 자식으로 코인 1개를 옮기는 데 1번의 이동이 추가됩니다. 따라서 전체 이동 횟수는 3번입니다.

풀이 접근 방법

이 문제는 재귀적 후위 순회(post-order traversal)를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 재귀 함수 solve()를 정의하고, 루트 노드를 인자로 전달합니다.
  • 노드가 null이면 0을 반환합니다.
  • 왼쪽 서브트리의 결과값 l과 오른쪽 서브트리의 결과값 r을 재귀적으로 계산합니다.
  • 각 서브트리에서 흘러 들어오거나 나가야 하는 코인 수의 절댓값(|l| + |r|)을 정답 변수 ans에 더합니다. 이 값이 곧 해당 노드를 기준으로 발생하는 이동 횟수입니다.
  • 현재 노드의 잉여 코인 수인 l + r + root.val - 1을 상위 호출로 반환합니다. 여기서 -1은 현재 노드가 자신을 위해 반드시 가져야 할 코인 하나를 의미합니다.
  • 메인 로직에서는 ans를 0으로 초기화한 뒤 solve(root)를 호출하고, 최종적으로 ans를 반환합니다.

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 ans;
   int solve(TreeNode* root){
      if(!root)return 0;
      int l = solve(root->left);
      int r = solve(root->right);
      ans += abs(l) + abs(r);
      return l + r + root->val - 1;
   }
   int distributeCoins(TreeNode* root) {
      ans = 0;
      solve(root);
      return ans;
   }
};
main(){
   vector<int> v = {0,3,0};
   TreeNode *root = make_tree(v);
   Solution ob;
   cout << (ob.distributeCoins(root));
}

입력

[0,3,0]

출력

3

복잡도 분석

이 알고리즘은 트리의 모든 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(N)입니다. 재귀 호출 스택의 깊이는 트리의 높이에 비례하므로, 공간 복잡도 역시 균형 잡힌 트리의 경우 O(log N), 편향된 트리의 경우 최악 O(N)입니다.