백준13516 트리와 쿼리 7

문제 링크

  • http://icpc.me/13516

사용 알고리즘

  • sqrt decomposition

풀이

intended solution은 Tree DP를 HLD로 최적화해주는 것인데, 트리와 쿼리 6에 비해 구현이 매우 더럽다는 소문을 들었습니다.
저는 그냥 sqrt decomposition을 사용했습니다.

쿼리를 $O(\sqrt Q)$개씩 묶어서 처리를 할 것입니다. $O(\sqrt Q)$개씩 분할한 각각의 조각을 “쿼리 묶음”이라고 합시다.
각 쿼리 묶음을 처리하기 전에, 해당 쿼리 묶음에서 색깔이나 가중치가 바뀌지 않으면서 색깔이 같은 인접한 정점들을 미리 union해줄 것입니다. 한 쿼리 묶음에서 색깔이나 가중치가 바뀌는 정점은 최대 $O(\sqrt Q)$개입니다.
정점들을 잘 묶어준 트리를 만들었으면, 쿼리 묶음에 있는 쿼리들을 하나씩 보면서 1/3번 쿼리는 그냥 색/가중치를 바꿔주면 되고, 2번 쿼리는 압축된 트리 위에서 bfs를 돌려주면 됩니다.

union-find에서 각 집합에 속한 정점의 가중치의 최댓값을 들고 있으면 간단하게 처리해줄 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#pragma GCC optimize ("O3")
#include <bits/stdc++.h>
#define x first
#define y second
using namespace std;

typedef pair<int, int> p;

int arr[101010];
struct UF{
    int par[101010];
    int mx[101010];
    void init(){
        for(int i=1; i<=100000; i++) par[i] = i, mx[i] = arr[i];
    }
    int find(int v) { return v == par[v] ? v : par[v] = find(par[v]); }
    void merge(int u, int v){
        u = find(u), v = find(v);
        if(u ^ v) par[u] = v, mx[v] = max(mx[v], mx[u]);
    }
    int siz(int v) { return mx[find(v)]; }
}uf;

int n;
int color[101010];
p qry[101010];
vector<int> inp[101010];
vector<int> g[101010];
int use[101010];
int upd[101010];

void bucket_init(int s, int e){
    memset(use, 0, sizeof use);
    uf.init();
    for(int i=s; i<=e; i++) if(qry[i].x != 2) use[qry[i].y] = 1;
    for(int i=1; i<=n; i++){
        g[i].clear();
        if(use[i]) continue;
        for(auto j : inp[i]) if(!use[j] && color[i] == color[j]) uf.merge(i, j);
    }
    for(int i=1; i<=n; i++){
        for(auto j : inp[i]){
            if(uf.find(i) != uf.find(j)) g[uf.find(i)].push_back(uf.find(j));
        }
    }
}

int chk[101010], ts;
queue<int> q;
int query(int v){
    ts++; int ans = 0;
    q.push(v); chk[v] = ts;
    while(q.size()){
        int now = q.front(); q.pop();
        ans = max(ans, uf.siz(now));
        for(auto i : g[now]){
            if(chk[i] == ts) continue;
            if(color[now] != color[i]) continue;
            chk[i] = ts; q.push(i);
        }
    }
    return ans;
}

void bucket(int s, int e){
    bucket_init(s, e);
    for(int i=s; i<=e; i++){
        if(qry[i].x == 1) color[qry[i].y] ^= 1;
        else if(qry[i].x == 2){
            cout << query(uf.find(qry[i].y)) << "\n";
        }else{
            arr[qry[i].y] = upd[i];
            uf.mx[qry[i].y] = upd[i];
        }
    }
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n;
    for(int i=1; i<n; i++){
        int s, e; cin >> s >> e;
        inp[s].push_back(e); inp[e].push_back(s);
    }
    for(int i=1; i<=n; i++) cin >> color[i];
    for(int i=1; i<=n; i++) cin >> arr[i];
    int q; cin >> q;
    for(int i=1; i<=q; i++){
        cin >> qry[i].x >> qry[i].y;
        if(qry[i].x == 3) cin >> upd[i];
    }

    for(int i=0; ; i++){
        int s = 400*i+1;
        int e = 400*(i+1);
        e = min(q, e);
        if(s > e) break;
        bucket(s, e);
    }
}