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

C++ 이진 트리 가지치기: 1을 포함하지 않는 서브트리 제거하기

문제 소개

모든 노드의 값이 0 또는 1로만 구성된 이진 트리(binary tree)가 있다고 가정해 봅시다. 이때 값이 1인 노드를 하나도 포함하지 않는 모든 서브트리(subtree)를 잘라내어(prune), 결과적으로 동일한 구조의 트리를 만드는 것이 목표입니다.

예를 들어 다음과 같은 트리가 주어졌다면,

C++ 이진 트리 가지치기: 1을 포함하지 않는 서브트리 제거하기

트리 말단에 있는 0으로만 이루어진 가지들이 모두 잘려 나간 트리가 최종 결과가 됩니다.

해결 전략: 재귀와 후위 순회

이 문제는 재귀(Recursion)를 활용하면 매우 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 후위 순회(postorder)처럼 자식 노드를 먼저 처리한 뒤 부모 노드를 판단하는 것입니다. 자식들이 먼저 정리되어야 현재 노드가 리프 노드가 되었는지, 그리고 그 값이 0인지 정확히 확인할 수 있기 때문입니다.

구체적인 알고리즘 단계는 다음과 같습니다.

  1. 재귀 메서드 solve(node)를 정의합니다.
  2. 노드가 NULL이면 NULL을 반환합니다. (베이스 케이스)
  3. node->left = solve(node의 왼쪽 자식)을 수행합니다.
  4. node->right = solve(node의 오른쪽 자식)을 수행합니다.
  5. 왼쪽 자식이 NULL이고, 오른쪽 자식도 NULL이며, 현재 노드의 값이 0이라면 NULL을 반환하여 해당 노드를 트리에서 잘라냅니다.
  6. 그 외의 경우에는 노드를 그대로 반환합니다.

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 inorder(TreeNode *root){
    if(root){
        inorder(root->left);
        cout << root->val << ", ";
        inorder(root->right);
    }
}
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(){
    TreeNode *root = new TreeNode(1);
    root->left = new TreeNode(1);
    root->right = new TreeNode(0);
    root->left->left = new TreeNode(1);
    root->left->right = new TreeNode(1);
    root->right->left = new TreeNode(0);
    root->right->right = new TreeNode(1);
    root->left->left->left = new TreeNode(0);
    Solution ob;
    inorder(ob.pruneTree(root));
}

입력

TreeNode *root = new TreeNode(1);
root->left = new TreeNode(1);
root->right = new TreeNode(0);
root->left->left = new TreeNode(1);
root->left->right = new TreeNode(1);
root->right->left = new TreeNode(0);
root->right->right = new TreeNode(1);
root->left->left->left = new TreeNode(0);

출력

1, 1, 1, 1, 0, 1,

코드 동작 원리

pruneTree 함수는 트리의 가장 깊은 곳, 즉 리프 노드부터 위로 거슬러 올라가며 동작합니다. 자식 노드들의 처리가 끝나면 현재 노드를 검사하는데, 만약 양쪽 자식이 모두 잘려나가 사라졌고(둘 다 NULL) 자신의 값도 0이라면, 이 노드는 더 이상 1을 포함하는 서브트리에 기여하지 않으므로 NULL을 반환해 제거합니다.

예제 트리에서 루트의 왼쪽 자식(1)의 왼쪽 자식(1) 아래에 있는 값 0의 리프 노드는 조건에 해당하므로 제거됩니다. 마찬가지로 오른쪽 서브트리에서 값이 0인 리프 노드 역시 잘려 나갑니다. 반면, 자식이 하나라도 남아 있거나 자신의 값이 1인 노드는 그대로 유지됩니다. 최종적으로 중위 순회(inorder) 결과가 1, 1, 1, 1, 0, 1,로 출력되는 것을 확인할 수 있습니다.

시간 및 공간 복잡도

  • 시간 복잡도: O(N) — 트리의 모든 노드를 정확히 한 번씩 방문합니다. (N은 전체 노드 수)
  • 공간 복잡도: O(H) — 재귀 호출 스택이 트리의 높이(H)만큼 사용됩니다. 편향된 트리의 최악의 경우 O(N), 균형 잡힌 트리라면 O(log N)입니다.

마무리

이진 트리 가지치기 문제는 재귀적 사고력을 기르기에 좋은 대표적인 트리 연습 문제입니다. 자식 노드의 결과를 먼저 계산한 뒤 부모 노드를 결정하는 후위 순회 패턴은 트리 구조 변형 문제 전반에 널리 활용되므로, 이번 기회에 확실히 익혀 두시길 권장합니다.