[BOJ 17469] 트리의 색깔과 쿼리

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

17469번: 트리의 색깔과 쿼리

Tag : Offline Query, DSU, Small To Large

 

풀이

간선을 제거하는 것은 일단 생각하지 말자. 어떠한 정점 \(A\)에서 갈 수 있는 점은 같은 컴포넌트에 있는 모든 정점이다. 컴포넌트들에 대해 서로 다른 색의 개수를 구하려면 일단 컴포넌트마다 std::set을 가지고 있어야할 듯하다.

이 때 컴포넌트는 DSU로 처리하면 되니 2번쿼리를 std::set::size를 통해 구할 수 있다.

 

간선을 제거하는 것은 어떻게 해야할까? 이전에 DSU를 통해 컴포넌트를 관리하겠다고 했는데 그렇다면 간선을 제거하는 것보다는 연결하는 것이 편할 것이다. 이제 모든 쿼리를 뒤집어서 보면 간선을 연결하는 문제로 바꿀 수 있다. 간선을 연결하면서 컴포넌트들의 std::set을 합치는 과정은 Small To Large를 통해 빠르게 수행할 수 있다. 

 

시간복잡도 : \(O(NlogN + Q)\)

 

전체 코드

#include <bits/stdc++.h>
#define pb push_back
#define fi first
#define se second
#define all(x) ((x).begin()), ((x).end())
#define endl '\n'
using namespace std;
using ll = long long;

int N, Q;
set<int> st[101010];
int pa[101010], parent[101010];

int find(int a){
    if(a == pa[a]) return a;
    else return pa[a] = find(pa[a]);
}

void merge(int a, int b){
    a = find(a), b = find(b);
    if(a == b) return;
    if(st[a].size() < st[b].size()) swap(a,b);
    for(auto &i : st[b]) st[a].insert(i);
    pa[b] = a;
}

int main(){
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    iota(pa+1,pa+1+100000,1);
    cin >> N >> Q;
    for(int i = 2; i<=N; i++) cin >> parent[i];
    for(int i = 1; i<=N; i++){
        int x; cin >> x;
        st[find(i)].insert(x);
    }
    vector<pair<int,int>> query(N+Q-1);
    for(auto &[i, j] : query) cin >> i >> j;
    reverse(all(query));
    stack<int> ans;
    for(auto &[i, j] : query){
        if(i == 1) merge(j, parent[j]);
        else ans.push(st[find(j)].size());
    }
    while(ans.size()){
        cout << ans.top() << endl;
        ans.pop();
    }
}