C++로 두 개의 이진 트리 병합하기

두 개의 이진 트리 병합이란?

두 개의 이진 트리가 있고, 그중 하나를 다른 트리 위에 겹쳐 올린다고 가정해 봅시다. 이때 일부 노드는 서로 겹치게 되고, 나머지 노드는 겹치지 않습니다. 우리의 목표는 이 두 트리를 하나의 새로운 이진 트리로 병합하는 것입니다.

병합 규칙은 다음과 같습니다.

• 두 노드가 겹치는 경우: 두 노드의 값을 더한 합을 병합된 노드의 새로운 값으로 사용합니다.

• 한쪽 노드만 존재하는 경우: 값이 있는(null이 아닌) 노드를 그대로 새 트리의 노드로 사용합니다.

예를 들어 아래와 같은 두 트리가 있다면 −

병합 결과는 다음과 같습니다 −

해결 접근 방법

이 문제는 재귀(Recursion)를 이용하면 간단하게 해결할 수 있습니다. 단계별로 살펴보면 다음과 같습니다 −

• 병합 함수 mergeTrees()는 두 개의 트리 노드 n1과 n2를 매개변수로 받습니다.

• n1이 null이고 n2가 null이 아니면 n2를 반환하고, 반대로 n2가 null이고 n1이 null이 아니면 n1을 반환합니다. 두 노드 모두 null이라면 null을 반환합니다.

• 두 노드가 모두 존재하면 n1의 값에 n2의 값을 더합니다. (n1->val += n2->val)

• n1의 왼쪽 자식은 mergeTrees(n1의 왼쪽 자식, n2의 왼쪽 자식)의 결과로 설정합니다.

• n1의 오른쪽 자식 역시 같은 방식으로 재귀적으로 병합합니다.

• 병합이 완료된 n1을 반환합니다.

이 방식은 각 노드를 정확히 한 번씩만 방문하므로, 시간 복잡도는 O(max(m, n))입니다. 여기서 m과 n은 각각 두 트리의 노드 수입니다.

예제 코드 (C++)

아래 구현을 통해 더 잘 이해해 보겠습니다 −

#include <bits/stdc++.h>
using namespace std;
class TreeNode{
   public:
      int val;
      TreeNode *left, *right;
      TreeNode(int data){
         val = data;
         left = 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){
         temp->left = new TreeNode(val);
         return;
      }
      else{
         q.push(temp->left);
      }
      if(!temp->right){
         temp->right = new TreeNode(val);
         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;
}
void tree_level_trav(TreeNode*root){
   if (root == NULL) return;
   cout << "[";
   queue<TreeNode *> q;
   TreeNode *curr;
   q.push(root);
   q.push(NULL);
   while (q.size() > 1) {
      curr = q.front();
      q.pop();
      if (curr == NULL){
         q.push(NULL);
      }
      else {
         if(curr->left)
            q.push(curr->left);
         if(curr->right)
            q.push(curr->right);
         if(curr->val == 0){
            cout << "null" << ", ";
         }
         else{
            cout << curr->val << ", ";
         }
      }
   }
   cout << "]"<<endl;
}
class Solution {
public:
   TreeNode* mergeTrees(TreeNode* n1, TreeNode* n2) {
      if(!n1 && n2){
         return n2;
      }
      else if(!n2 && n1)return n1;
      else if(!n1 && !n2)return NULL;
      n1->val+=n2->val;
      n1->left = mergeTrees(n1->left,n2->left);
      n1->right = mergeTrees(n1->right,n2->right);
      return n1;
   }
};
main(){
   Solution ob;
   vector<int> v1 = {1,3,2,5};
   vector<int> v2 = {2,1,3,NULL,4,NULL,7};
   TreeNode *root1 = make_tree(v1);
   TreeNode *root2 = make_tree(v2);
   root1 = ob.mergeTrees(root1, root2);
   tree_level_trav(root1);
}

입력

[1,3,2,5]
[2,1,3,null,4,null,7]

출력

[3, 4, 5, 5, 4, null, 7]