(17469) 트리의 색깔과 쿼리

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

단계별로 풀어보기

43단계(유니온 파인드 2) 3번째

 

 

 

Small to Large Trick의 뭉쳐야 산다와 오프라인 쿼리의 트리 문제를 합친 문제

 

 

 

#include <iostream>
#include <vector>
#include <algorithm>
#include <set>
using namespace std;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  int n, q;
  
  cin >> n >> q;
  
  vector<int> uf(n+1), parent(n+1), query(n+q-1), ans;
  vector<set<int>> v(n+1);
  
  for(int i=2;i<=n;i++)  cin >> parent[i];
  
  for(int i=1;i<=n;i++){
    int c;
    
    cin >> c;
    
    v[i].insert(c);
  }
  
  for(int i=0;i<n+q-1;i++){
    int qtype, a;
    
    cin >> qtype >> a;
    
    query[i]= qtype==1 ? a : -a;
  }
  
  reverse(query.begin(), query.end());
  
  for(int i=1;i<=n;i++)  uf[i]= i;
  
  auto find= [&](auto self, int x) -> int {
    if(uf[x]==x)  return x;
    else  return uf[x]= self(self, uf[x]);
  };
  
  auto merge= [&](int x, int y){
    x= find(find, x);
    y= find(find, y);
    
    uf[y]= x;
  };
  
  for(auto now:query){
    if(now>0){
      int nowparent= find(find, parent[now]);
      
      if(v[nowparent].size()<v[now].size())  swap(v[nowparent], v[now]);
      
      v[nowparent].insert(v[now].begin(), v[now].end());
      v[now].clear();
      
      merge(nowparent, now);
    }else{
      ans.push_back(v[find(find, -now)].size());
    }
  }
  
  reverse(ans.begin(), ans.end());
  
  for(int res:ans)  cout << res << "\n";
  
  return 0;
}