문제 개요
n개의 노드로 구성된 이진 트리가 주어졌을 때, 원본 트리에서 정확히 하나의 간선을 제거하여 트리를 두 개로 분할했을 때, 두 트리의 노드 값 합이 서로 같아질 수 있는지 확인하는 것이 이번 문제의 목표입니다.
예를 들어 아래와 같은 트리가 입력으로 주어진다면,
출력 결과는 true가 됩니다.
해결 접근 방법
이 문제는 재귀적 깊이 우선 탐색(DFS)과 스택을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.
- 각 노드를 루트로 하는 부분 트리(subtree)의 합을 모두 계산하여 스택에 저장합니다.
- 전체 트리의 총합(totalSum)을 구한 후, 스택에 남아 있는 각 부분 트리의 합 x에 대해 x == totalSum − x가 성립하는지 확인합니다.
- 조건을 만족하는 부분 트리가 존재한다면, 해당 부분 트리와 나머지 부분을 연결하는 간선 하나만 잘라내면 두 트리의 합이 같아지므로 true를 반환합니다.
여기서 스택의 최상단 값은 루트 노드 전체의 합이므로 먼저 제거합니다. 이렇게 하면 간선을 실제로 잘랐을 때 만들어질 수 있는 부분 트리의 합만 검사하게 되어, 루트 단독으로 분리되는 잘못된 경우를 자연스럽게 배제할 수 있습니다.
알고리즘 단계
- 정수를 저장할 스택 st를 하나 선언합니다.
- 노드를 매개변수로 받는 solve() 함수를 정의합니다.
- 노드가 null이면 0을 반환합니다.
- leftSum := solve(노드의 왼쪽 자식)을 통해 왼쪽 부분 트리의 합을 구합니다.
- rightSum := solve(노드의 오른쪽 자식)을 통해 오른쪽 부분 트리의 합을 구합니다.
- curr := 노드의 값 + leftSum + rightSum 으로 현재 부분 트리의 합을 계산합니다.
- curr을 스택 st에 삽입(push)한 뒤 curr을 반환합니다.
- 메인 함수에서는 solve(root)를 호출하여 모든 부분 트리의 합을 스택에 쌓습니다.
- totalSum := 스택의 최상단 값을 꺼낸 뒤, 해당 요소를 스택에서 제거(pop)합니다. 이 값이 전체 트리의 합입니다.
- 스택이 빌 때까지 다음을 반복합니다.
- x := 스택의 최상단 값을 꺼내고(pop), y := totalSum − x를 계산합니다.
- x와 y가 같다면 true를 반환합니다.
- 반복이 끝날 때까지 조건을 만족하지 못했다면 false를 반환합니다.
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;
}
};
class Solution {
public:
stack <int> st;
int solve(TreeNode* node){
if (!node)
return 0;
int leftSum = solve(node->left);
int rightSum = solve(node->right);
int curr = node->val + leftSum + rightSum;
st.push(curr);
return curr;
}
bool checkEqualTree(TreeNode* root) {
solve(root);
int totalSum = st.top();
st.pop();
while (!st.empty()) {
int x = st.top();
st.pop();
int y = totalSum - x;
if (x == y)
return true;
}
return false;
}
};
main(){
Solution ob;
TreeNode *root = new TreeNode(5);
root->left = new TreeNode(10);
root->right = new TreeNode(10);
root->right->left = new TreeNode(2);
root->right->right = new TreeNode(3);
cout<<(ob.checkEqualTree(root));
}입력
TreeNode *root = new TreeNode(5); root->left = new TreeNode(10); root->right = new TreeNode(10); root->right->left = new TreeNode(2); root->right->right = new TreeNode(3);
출력
1
동작 원리 정리
위 예제에서 전체 트리의 합은 5 + 10 + 10 + 2 + 3 = 30입니다. 루트의 왼쪽 부분 트리 합은 10이고, 나머지 부분의 합은 30 − 10 = 20입니다. 그러나 오른쪽 부분 트리(10 + 2 + 3 = 15)를 분리하면 15 vs 15로 두 트리의 합이 정확히 같아집니다. 따라서 함수는 true(1)를 반환합니다.
이 알고리즘은 모든 노드를 한 번씩 방문하므로 시간 복잡도는 O(n), 스택에 부분 트리의 합을 저장하므로 공간 복잡도 역시 O(n)입니다.