C++로 해결하는 두 개의 BST에서 Two Sum 찾기

문제 설명

두 개의 이진 탐색 트리(Binary Search Tree, BST)가 주어졌을 때, 첫 번째 트리의 어떤 노드와 두 번째 트리의 어떤 노드를 골라 그 값의 합이 주어진 정수 target과 일치한다면 true를 반환하고, 그런 경우가 존재하지 않으면 false를 반환하는 문제입니다.

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

첫 번째 트리는 [2,1,4], 두 번째 트리는 [1,0,3]으로 구성되어 있고, target이 5라면 결과는 true입니다. 실제로 첫 번째 트리의 노드 2와 두 번째 트리의 노드 3을 더하면 5가 되기 때문입니다.

해결 접근 방법

이 문제는 해시 맵(hash map)을 활용해 효율적으로 해결할 수 있습니다. 핵심 아이디어는 첫 번째 트리의 모든 노드 값을 맵에 기록해 둔 뒤, 두 번째 트리를 순회하면서 'target - 현재 값'이 맵에 존재하는지 확인하는 것입니다.

구체적인 단계는 다음과 같습니다.

값과 해당 값이 속한 트리의 번호를 저장할 맵 s를 정의합니다.
노드(node), 목표값(target), 트리 번호(nodeNumber)를 매개변수로 받는 check() 메서드를 정의합니다.
노드가 유효하지 않으면(null이면) false를 반환합니다.
curr := 현재 노드의 값, req := target - curr 로 설정합니다.
req가 맵 s에 존재하고, s[req]가 현재 nodeNumber와 다르면(즉, 서로 다른 트리에서 온 값이라면) true를 반환합니다.
s[curr] := nodeNumber 로 현재 값을 맵에 기록합니다.
왼쪽 자식과 오른쪽 자식에 대해 check()를 재귀적으로 호출한 결과를 OR 연산으로 반환합니다.
메인 메서드(twoSumBSTs)에서는 먼저 flag := check(root1, target, 1)로 첫 번째 트리를 순회한 후, return check(root2, target, 2)로 두 번째 트리를 순회하며 최종 결과를 반환합니다.

여기서 nodeNumber를 함께 저장하는 이유는, 같은 트리 내부의 두 노드 합이 target과 일치하는 경우를 잘못 판정하지 않도록 하기 위함입니다. 즉, 반드시 서로 다른 트리에 속한 노드들의 합만 인정됩니다.

예제 코드 (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:
   map <int,int> s;
   bool check(TreeNode* node, int target,int nodeNumber){
      if(!node)return false;
      int curr = node->val;
      int req = target - curr;
      if(s.find(req)!=s.end() && s[req]!=nodeNumber)return true;
      s[curr]=nodeNumber;
      return check(node->left,target,nodeNumber) || check(node->right,target,nodeNumber);
   }
   bool twoSumBSTs(TreeNode* root1, TreeNode* root2, int target) {
      bool flag = check(root1,target,1);
      return check(root2,target,2);
   }
};
main(){
   vector<int> v1 = {2,1,4};
   vector<int> v2 = {1,0,3};
   TreeNode *r1 = make_tree(v1);
   TreeNode *r2 = make_tree(v2);
   Solution ob;
   cout <<ob.twoSumBSTs(r1, r2, 5);
}

입력

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

출력

복잡도 분석

이 알고리즘은 각 트리의 모든 노드를 한 번씩 방문하므로 시간 복잡도는 O(N + M)입니다(N, M은 각각 두 트리의 노드 수). 또한 맵에 각 트리의 노드 값을 저장해야 하므로 공간 복잡도 역시 O(N + M)입니다. 만약 두 트리가 균형 잡힌 BST라면, 중위 순회(inorder traversal)로 각 트리를 정렬된 배열로 변환한 뒤 투 포인터(two pointer) 기법을 사용하는 방법도 고려할 수 있습니다.