문제 개요
이진 트리의 루트(root) 노드가 주어지고, 트리를 구성하는 모든 노드의 값은 0 또는 1이라고 가정해 봅시다. 우리가 구해야 하는 것은 값이 1인 노드를 하나도 포함하지 않는 서브트리(subtree)를 모두 잘라낸(pruning) 결과 트리입니다.
예를 들어 트리가 다음과 같다고 하면 −
해결 접근 방법
이 문제는 후위 순회(post-order traversal) 기반의 재귀 함수를 사용하면 깔끔하게 해결할 수 있습니다. 자식 노드를 먼저 처리한 뒤 부모 노드를 판단해야, 리프 노드부터 차례대로 불필요한 가지를 잘라낼 수 있기 때문입니다.
구체적인 알고리즘은 다음과 같습니다.
- 노드를 인자로 받는 재귀 메서드 solve()를 정의합니다.
- 현재 노드가 null이면 null을 그대로 반환합니다.
- 왼쪽 자식에 대해 solve(node->left)를 호출한 결과를 다시 node->left에 저장합니다.
- 오른쪽 자식에 대해 solve(node->right)를 호출한 결과를 다시 node->right에 저장합니다.
- 왼쪽 자식과 오른쪽 자식이 모두 null이고 현재 노드의 값이 0이라면, 이 노드는 1을 포함하는 서브트리에 기여하지 않으므로 null을 반환해 제거합니다.
- 그 외의 경우에는 현재 노드를 그대로 반환합니다.
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){
temp->left = new TreeNode(val);
return;
}else{
q.push(temp->left);
}
if(!temp->right){
temp->right = new TreeNode(val);
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;
}
void tree_level_trav(TreeNode*root){
if (root == NULL) return;
cout << "[";
queue<TreeNode *> q;
TreeNode *curr;
q.push(root);
q.push(NULL);
while (q.size() > 1) {
curr = q.front();
q.pop();
if (curr == NULL){
q.push(NULL);
} else {
if(curr->left)
q.push(curr->left);
if(curr->right)
q.push(curr->right);
if(curr == NULL){
cout << "null" << ", ";
}else{
cout << curr->val << ", ";
}
}
}
cout << "]"<<endl;
}
class Solution {
public:
TreeNode* pruneTree(TreeNode* node) {
if(!node)return NULL;
node->left = pruneTree(node->left);
node->right = pruneTree(node->right);
if(!node->left && !node->right && !node->val){
return NULL;
}
return node;
}
};
main(){
vector<int> v = {1,1,0,1,1,0,1,0};
TreeNode *root = make_tree(v);
Solution ob;
tree_level_trav(ob.pruneTree(root));
}
입력
[1,1,0,1,1,0,1,0]
출력
[1, 1, 0, 1, 1, 1]
입력 트리에서 값이 0인 리프 노드들이 제거되고, 그로 인해 1을 포함하게 된 부모 노드들은 그대로 유지된 것을 확인할 수 있습니다.
복잡도 분석
시간 복잡도: O(n) — 트리의 모든 노드를 정확히 한 번씩 방문하므로 노드 개수 n에 비례합니다.
공간 복잡도: O(h) — 재귀 호출 스택의 최대 깊이는 트리의 높이 h에 의해 결정됩니다. 균형 잡힌 트리라면 O(log n), 편향된 트리라면 최악의 경우 O(n)이 됩니다.