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

C++로 구현하는 이진 트리에서 최대 합 BST 찾기

문제 개요

이진 트리의 루트(root)가 주어졌을 때, 그중에서 이진 탐색 트리(BST) 조건을 동시에 만족하는 서브트리를 골라 해당 서브트리에 속한 모든 노드 값의 합이 최대가 되도록 구하는 문제입니다.

예를 들어 다음과 같은 트리가 입력으로 주어진다면,

C++로 구현하는 이진 트리에서 최대 합 BST 찾기

정답은 20이 됩니다. 선택된 BST에 포함된 모든 노드 값의 합이 20이기 때문입니다.

접근 방법

이 문제는 후위 순회(postorder traversal) 기반의 재귀적 풀이로 깔끔하게 해결할 수 있습니다. 각 노드를 기준으로 하는 서브트리가 BST인지 여부, 노드 개수, 서브트리 내 최댓값과 최솟값, 노드 값의 합을 자식 노드에서 부모 노드 방향으로 전달하며 처리하는 방식입니다.

핵심 아이디어: Data 구조체

재귀 호출 결과를 담기 위해 sz(노드 개수), maxVal(서브트리 최댓값), minVal(서브트리 최솟값), ok(BST 여부), sum(노드 값의 합) 멤버를 가지는 Data 구조체를 사용합니다.

알고리즘 단계

  1. sz, maxVal, minVal, ok, sum 멤버를 가지는 Data 구조체를 정의합니다. 생성자는 (sz, minVal, maxVal, ok) 순서로 값을 받고, sum은 0으로 초기화합니다.
  2. 전역 결과 변수 ret := 0으로 초기화합니다.
  3. 트리 노드를 인자로 받는 solve() 함수를 정의합니다.
  4. 노드가 null이거나 노드 값이 0이면, 빈 서브트리를 의미하는 Data(0, INF, -INF, true) 객체를 반환합니다.
  5. left := solve(노드의 왼쪽 자식)
  6. right := solve(노드의 오른쪽 자식)
  7. Data 타입 객체 curr을 생성하고 curr.ok := false로 설정합니다.
  8. node->val >= right.minVal이라면 BST 성질을 위반하므로 curr을 즉시 반환합니다.
  9. node->val <= left.maxVal이라면 마찬가지로 curr을 반환합니다.
  10. left.ok와 right.ok가 모두 true라면 현재 서브트리는 유효한 BST이므로 다음을 수행합니다.
    • curr.sum = node->val + left.sum + right.sum
    • ret = max(curr.sum, ret)
    • curr.sz = 1 + left.sz + right.sz
    • curr.ok = true
    • curr.maxVal = max(node->val, right.maxVal)
    • curr.minVal = min(node->val, left.minVal)
  11. curr을 반환합니다.
  12. 메인 함수에서는 ret := 0으로 초기화한 뒤 solve(root)를 호출하고, 최종적으로 ret을 반환합니다.

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;
}
struct Data{
   int sz;
   int maxVal;
   int minVal;
   bool ok;
   int sum;
   Data(){}
   Data(int a, int b, int c, bool d){
      sz = a;
      minVal = b;
      maxVal = c;
      ok = d;
      sum = 0;
   }
};
class Solution {
   public:
   int ret = 0;
   Data solve(TreeNode* node){
      if (!node || node->val == 0)
      return Data(0, INT_MAX, INT_MIN, true);
      Data left = solve(node->left);
      Data right = solve(node->right);
      Data curr;
      curr.ok = false;
      if (node->val >= right.minVal) {
         return curr;
      }
      if (node->val <= left.maxVal) {
         return curr;
      }
      if (left.ok && right.ok) {
         curr.sum = node->val + left.sum + right.sum;
         ret = max(curr.sum, ret);
         curr.sz = 1 + left.sz + right.sz;
         curr.ok = true;
         curr.maxVal = max(node->val, right.maxVal);
         curr.minVal = min(node->val, left.minVal);
      }
      return curr;
   }
   int maxSumBST(TreeNode* root){
      ret = 0;
      solve(root);
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v =
   {1,4,3,2,4,2,5,NULL,NULL,NULL,NULL,NULL,NULL,4,6};
   TreeNode *root = make_tree(v);
   cout << (ob.maxSumBST(root));
}

입력

{1,4,3,2,4,2,5,NULL,NULL,NULL,NULL,NULL,NULL,4,6}

출력

20

복잡도 분석

각 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(N)이며, 재귀 호출 스택의 깊이가 트리의 높이에 비례하므로 공간 복잡도는 O(H)(H는 트리의 높이)입니다. 편향된 트리의 경우 최악에 O(N)까지 늘어날 수 있습니다.