BOJ 25549 - 트리의 MEX [C++]

문제

https://www.acmicpc.net/problem/25549

 

$N$개의 정점으로 구성된 루트 있는 트리의 각 정점에 정수가 하나씩 쓰여있다.

$mex(i)$를 $i$번 정점을 루트로 하는 서브트리의 정점에 적혀 있지 않은 수 중에서 가장 작은 음이 아닌 정수라고 할 때, $1 \le i \le N$인 모든 $i$에 대해 $mex(i)$를 출력해라.

 

풀이

정점 $i$의 서브트리에 적혀 있는 모든 수의 집합(std::set)을 $S_i$라 하자.

dfs를 돌면서 부모 정점 $n$과 그 자식 정점들 $c_1, c_2, \ldots, c_j$에 대해, $S_n$과 $S_{c_1}, S_{c_2}, \ldots, S_{c_j}$를 모두 합쳐서 $mex(n)$을 직접 구하는 풀이가 가능하다.

집합을 small to large trick으로 합치면 시간복잡도가 $O(N\log{N})$이고, $S_n$에 원소가 삽입될 때마다 $mex(n)$은 절대 줄어들 수 없기 때문에 이 또한 결국 $O(N\log{N})$정도로 떨어질 것이다.

 

#include <bits/stdc++.h>
using namespace std;

set<int> S[200001]{};
int mex[200001]{}, N, R;
vector<vector<int>> V(200001);
int ans[200001]{};

void dfs(int n, int p){
    for(int i:V[n]){
        dfs(i,n);
        if(S[n].size() < S[i].size())   swap(S[n], S[i]), swap(mex[n], mex[i]);
        for(int j:S[i]){
            if(S[n].count(j))   continue;
            S[n].insert(j);
            if(j != mex[n]) continue;
            int new_mex = j;
            auto k = S[n].find(j);
            while(k != S[n].end() && *k == new_mex) new_mex++, k++;
            mex[n] = new_mex;
        }
    }
    ans[n] = mex[n];
}

int main(){
    cin.tie(0)->sync_with_stdio(0);

    cin>>N;
    for(int i=1;i<=N;i++){
        int a;
        cin>>a;
        if(a == -1) {R = i;continue;}
        V[a].push_back(i);
    }
    for(int i=1;i<=N;i++){
        int a;
        cin>>a;
        S[i].insert(a);
        mex[i] = a == 0 ? 1 : 0;
    }
    dfs(R, 0);
    for(int i=1;i<=N;i++)   cout<<ans[i]<<'\n';

}