C++로 이진 트리의 유니밸류(단일값) 서브트리 개수 구하기

문제 개요

하나의 이진 트리가 주어졌을 때, 유니밸류 서브트리(uni-value subtree)의 개수를 세는 것이 이번 문제의 목표입니다. 여기서 유니밸류 서브트리란 해당 서브트리에 속한 모든 노드의 값이 동일한 서브트리를 의미합니다.

예를 들어 트리가 다음과 같이 주어진 경우를 살펴보겠습니다.

root = [5,1,5,5,5,null,5]

위 트리에서 값이 모두 같은 서브트리는 총 4개이므로, 출력 결과는 4가 됩니다.

해결 접근 방법

이 문제는 재귀적 후위 순회(post-order traversal)를 이용하면 깔끔하게 해결할 수 있습니다. 각 노드마다 "내 서브트리가 유니밸류인가?"를 판단하고, 참일 경우 카운트를 1씩 늘려가는 방식입니다.

핵심 로직은 다음과 같습니다.

solve() 함수를 정의하고 노드를 인자로 전달합니다.
노드가 비어 있다면 true를 반환합니다. (빈 트리는 항상 유니밸류로 간주)
left = solve(노드의 왼쪽 자식), right = solve(노드의 오른쪽 자식)을 호출합니다.
left 또는 right 중 하나라도 false라면 false를 반환합니다.
왼쪽 자식이 존재하는데 부모 노드와 값이 다르면 false를 반환합니다.
오른쪽 자식이 존재하는데 부모 노드와 값이 다르면 false를 반환합니다.
모든 조건을 통과했다면 카운터(ret)를 1 증가시키고 true를 반환합니다.

메인 함수에서는 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;
}
class Solution {
public:
    int ret;
    bool solve(TreeNode* node){
        if (!node || node->val == 0)
            return true;
        bool left = solve(node->left);
        bool right = solve(node->right);
        if (!left || !right)
            return false;
        if (node->left && node->left->val != 0 && node->val != node->left->val)
            return false;
        if (node->right && node->right->val != 0 && node->val != node->right->val)
            return false;
        ret++;
        return true;
    }
    int countUnivalSubtrees(TreeNode* root){
        ret = 0;
        solve(root);
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {5,1,5,5,5,NULL,5};
    TreeNode *root = make_tree(v);
    cout << (ob.countUnivalSubtrees(root));
}

입력

{5,1,5,5,5,NULL,5}

출력

복잡도 분석

이 알고리즘은 트리의 모든 노드를 정확히 한 번씩만 방문하므로, 시간 복잡도는 O(N)(N은 노드의 개수)입니다. 재귀 호출에 따른 스택 공간이 필요하므로 공간 복잡도는 최악의 경우 편향된 트리일 때 O(N), 균형 잡힌 트리일 때 O(log N)입니다.