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

C++로 오염된 이진 트리 복구 후 요소 찾기

문제 설명

다음과 같은 규칙을 따르는 이진 트리가 있다고 가정해 봅시다.

  • 루트 노드의 값은 항상 root.val == 0 입니다.

  • treeNode.val이 x이고 왼쪽 자식이 null이 아니라면, treeNode.left.val = 2 * x + 1 입니다.

  • treeNode.val이 x이고 오른쪽 자식이 null이 아니라면, treeNode.right.val = 2 * x + 2 입니다.

그런데 이 이진 트리가 '오염(contaminated)'되었다고 합시다. 즉, 트리의 모든 노드 값이 -1로 변경된 상태입니다. 우리는 먼저 이진 트리를 원래 상태로 복구한 뒤, 다음 두 기능을 수행하는 FindElements 클래스를 구현해야 합니다.

  • FindElements(TreeNode* root): 오염된 이진 트리로 객체를 초기화하며, 내부적으로 트리를 먼저 복구합니다.

  • bool find(int target): 복구된 이진 트리에 target 값이 존재하면 true, 존재하지 않으면 false를 반환합니다.

예를 들어 트리가 다음과 같이 주어졌다고 해보겠습니다.

C++로 오염된 이진 트리 복구 후 요소 찾기

트리를 복구한 후 값 1, 3, 5를 차례대로 찾는다면 결과는 각각 true, true, false가 됩니다. 1과 3은 복구된 트리에 존재하지만, 5는 존재하지 않기 때문입니다.

접근 방법

이 문제는 깊이 우선 탐색(DFS)과 집합(set)을 활용하면 깔끔하게 해결할 수 있습니다. 단계별로 살펴보겠습니다.

  • 정수 값을 저장할 집합 a를 정의합니다.

  • dfs() 메서드를 정의합니다. 이 메서드는 root와 rootVal을 매개변수로 받으며, rootVal의 초기값은 -1입니다.

  • root가 null이면 그대로 반환합니다.

  • rootVal이 -1이라면 현재 노드의 값을 0으로 설정하고, 그렇지 않다면 rootVal로 설정합니다.

  • 현재 노드의 값을 집합 a에 삽입합니다.

  • 왼쪽 자식에는 dfs(left, 2 * 노드값 + 1), 오른쪽 자식에는 dfs(right, 2 * 노드값 + 2)를 재귀적으로 호출합니다.

  • 객체 초기화(트리 복구) 시에는 dfs(root, -1)을 호출합니다.

  • 요소를 찾을 때는 해당 값이 집합 a에 존재하는지만 확인하면 됩니다. 존재하면 true, 아니면 false를 반환합니다.

핵심 아이디어는 트리의 구조적 규칙(자식 노드 값 = 부모 노드 값 × 2 + 1 또는 + 2)을 이용해 오염된 값을 역으로 복원한다는 점입니다. 복구 과정에서 모든 유효한 값을 집합에 미리 저장해 두면, 이후 찾기 연산은 단순한 조회만으로 처리할 수 있습니다.

구현 예제

아래 코드를 통해 더 잘 이해해 보겠습니다.

#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 FindElements {
   public:
   set <int> a;
   void dfs(TreeNode* root,int rootVal=-1){
      if(!root)return;
      root->val = rootVal == -1?0:rootVal;
      a.insert(root->val);
      dfs(root->left,2*root->val + 1);
      dfs(root->right, 2*root->val + 2);
   }
   FindElements(TreeNode* root) {
      dfs(root);
   }
   bool find(int t) {
      return a.find(t)!=a.end();
   }
};
main(){
   vector<int> v = {-1,-1,-1,-1,-1};
   TreeNode *root = make_tree(v);
   FindElements ob(root);
   cout << (ob.find(1)) << endl;
   cout << (ob.find(3)) << endl;
   cout << (ob.find(5)) << endl;
}

입력

Initialize the tree with [-1,-1,-1,-1,-1], then call find(1), find(3) and find(5)

출력

1
1
0

출력 결과를 보면 find(1)과 find(3)은 1(true)을 반환하고, find(5)는 0(false)을 반환합니다. 이는 복구된 트리에 1과 3은 존재하지만 5는 존재하지 않음을 의미합니다.

복잡도 분석

  • 시간 복잡도: 트리 복구(DFS 순회)는 모든 노드를 한 번씩 방문하므로 O(n)이며, find 연산은 std::set(균형 이진 탐색 트리)에서의 조회이므로 O(log n)입니다.

  • 공간 복잡도: 복구된 값을 저장하는 집합과 재귀 호출 스택으로 인해 O(n)의 추가 공간이 필요합니다.

이처럼 DFS로 한 번만 전체 트리를 순회하며 값을 미리 수집해 두면, 이후 반복되는 찾기 질의를 매우 빠르게 처리할 수 있다는 것이 이 풀이의 장점입니다.