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

C++로 두 이진 탐색 트리의 모든 요소를 오름차순으로 병합하기

문제 개요

두 개의 이진 탐색 트리(Binary Search Tree, BST)가 주어졌을 때, 두 트리에 존재하는 모든 요소를 하나의 리스트에 담아 오름차순으로 반환하는 것이 이번 문제의 목표입니다.

예를 들어 아래와 같은 두 트리가 있다고 가정해 보겠습니다.

C++로 두 이진 탐색 트리의 모든 요소를 오름차순으로 병합하기

이 경우 출력 결과는 [0, 1, 1, 2, 3, 4]가 됩니다.

해결 접근 방식

이진 탐색 트리를 중위 순회(In-order Traversal)하면 값이 항상 오름차순으로 나온다는 성질을 활용하면 이 문제를 효율적으로 풀 수 있습니다. 재귀 호출 대신 스택 두 개를 사용해 반복적으로 순회하면서, 마치 병합 정렬(merge sort)에서 두 개의 정렬된 배열을 합치듯 두 스트림을 하나로 병합하는 방식입니다.

구체적인 알고리즘 단계는 다음과 같습니다.

  • 결과를 저장할 배열 ans와 두 개의 스택 st1, st2를 선언합니다.
  • curr1 := root1, curr2 := root2로 초기화합니다.
  • root1 노드와 그 왼쪽 자식들을 모두 st1에 push하고, root2 노드와 그 왼쪽 자식들을 모두 st2에 push합니다.
  • st1 또는 st2 중 하나라도 비어 있지 않은 동안 다음을 반복합니다.
    • st1이 비어 있지 않고, (st2가 비어 있거나 st1의 top 값이 st2의 top 값보다 작거나 같은 경우)
      • temp := st1의 top, st1에서 pop
      • temp의 값을 ans에 추가
      • temp의 오른쪽 서브트리와 그 왼쪽 자식들을 st1에 push
    • 그렇지 않은 경우
      • temp := st2의 top, st2에서 pop
      • temp의 값을 ans에 추가
      • temp의 오른쪽 서브트리와 그 왼쪽 자식들을 st2에 push
  • ans를 반환합니다.

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:
   void pushLeft(stack <TreeNode*>& st, TreeNode* root){
      TreeNode* curr = root;
      while(curr){
         st.push(curr);
         curr = curr->left;
      }
   }
   vector<int> getAllElements(TreeNode* root1, TreeNode* root2) {
      vector <int> ans;
      stack <TreeNode*> st1, st2;
      TreeNode* curr1 = root1;
      TreeNode* curr2 = root2;
      pushLeft(st1, curr1);
      pushLeft(st2, curr2);
      while(!st1.empty() || !st2.empty()){
         TreeNode* temp;
         if(!st1.empty() && (st2.empty() || st1.top()->val <= st2.top()->val)){
            temp = st1.top();
            st1.pop();
            ans.push_back(temp->val);
            pushLeft(st1, temp->right);
         }
         else{
            temp = st2.top();
            st2.pop();
            ans.push_back(temp->val);
            pushLeft(st2, temp->right);
         }
      }
      return ans;
   }
};
main(){
   vector<int> v = {2,1,4};
   TreeNode *root1 = make_tree(v);
   v = {1,0,3};
   TreeNode *root2 = make_tree(v);
   Solution ob;
   print_vector(ob.getAllElements(root1, root2));
}

입력

[2,1,4]
[1,0,3]

출력

[0,1,1,2,3,4]

복잡도 분석

시간 복잡도: O(N + M). N과 M은 각각 두 트리의 노드 개수입니다. 모든 노드를 정확히 한 번씩 방문하므로 전체 노드 수에 비례하는 시간이 걸립니다.

공간 복잡도: O(H1 + H2). H1과 H2는 각 트리의 높이입니다. 스택에는 어느 시점이든 각 트리의 높이만큼의 노드만 저장되며, 결과 배열을 제외한 추가 공간은 트리 높이에 의해 결정됩니다.