Statement

给一个树，nn 个点，有点权，初始根是 11。mm 个操作，种类如下：

  • 1 x\texttt{1}~x：将树根换为 xx。
  • 2 x y\texttt{2}~x~y：给出两个点 x,yx,y，从 xx 的子树中选每一个点，yy 的子树中选每一个点，求点权相等的情况数。

数据范围：1≤n≤1051\le n\le10^5，1≤m≤5×1051\le m\le5\times10^5 , 1≤ai≤1091\le a_i\le10^9。

Solution

根号做法，按颜色的出现次数根号分治。

对于出现次数 >n>\sqrt n 的颜色

最多 n\sqrt n 种颜色，那么对于每种颜色 O(n)O(n) 地统计 xx 子树内出现次数 cntx\text{cnt}_x。询问 x,yx,y 的答案就是 cntx⋅cnty\text{cnt}_x\cdot\text{cnt}_y。

当然由于换根，可能 xx 的答案变成了 sumx−cntprex,root\text{sum}_x-\text{cnt}_{\text{pre}_{\text{x,root}}}。

对于出现次数 <n<\sqrt n 的颜色

可以把所有同色点对拎出来，记为 (u,v)(u,v)。那么一个询问 x,yx,y 就是要求 uu 在 xx 子树内且 vv 在 yy 子树内的点对 (u,v)(u,v) 的个数。二维数点问题，可以带 log⁡\log 完成。

注意到修改次数 O(nn)O(n\sqrt n) 而查询次数 O(m)O(m)，所以把树状数组换成修改 O(1)O(1) 询问 O(n)O(\sqrt n) 的分块可以做到 O((n+m)n)O((n+m)\sqrt n)。

Code

其实数据范围对这个做法不友好，但是 shadowice1984 都能过！！1

可能是我写得太丑了吧 ╥﹏╥…

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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
#include <bits/stdc++.h>
namespace INPUT_SPACE {
const int BS = (1 << 24) + 5;
char Buffer[BS], *HD, *TL;
char gc() {
    if (HD == TL) TL = (HD = Buffer) + fread(Buffer, 1, BS, stdin);
    return (HD == TL) ? EOF : *HD++;
}
void input(int& t) {
    int x, ch, sgn = 1;
    while (((ch = gc()) < '0' || ch > '9') && ch != '-')
        ;
    if (ch == '-')
        sgn = -1, x = 0;
    else
        x = ch ^ '0';
    while ((ch = gc()) >= '0' && ch <= '9') x = (x << 1) + (x << 3) + (ch ^ '0');
    t = x * sgn;
}
}  // namespace INPUT_SPACE
using INPUT_SPACE::input;
typedef long long LL;
const int N = 1e5 + 3, M = 5e5 + 3, SN = 503, LogN = 18;
int n, q, a[N], B, in[N], out[N], pcolor = 0, color[SN], cnt[N][SN];
std::vector<int> upd[N];
namespace Gragh {
struct Edge {
    int to, next;
} edges[M];
int head[N], pedges = 0;
void addEdges(int u, int v) {
    edges[++pedges] = {v, head[u]}, head[u] = pedges;
    edges[++pedges] = {u, head[v]}, head[v] = pedges;
}
int pdfn = 0, f[N][LogN + 1], dep[N];
std::vector<int> vis[N];
void dfs(int u) {
    for (int k = 1; k <= LogN; ++k) f[u][k] = f[f[u][k - 1]][k - 1];
    in[u] = ++pdfn;
    bool tag = 0;
    for (int i = 1; i <= pcolor; ++i)
        if (color[i] == a[u]) {
            ++cnt[u][i], tag = 1;
            break;
        }
    if (!tag) {
        vis[a[u]].emplace_back(u);
        for (int i : vis[a[u]]) {
            upd[in[u]].emplace_back(in[i]);
            if (u != i) upd[in[i]].emplace_back(in[u]);
        }
    }
    for (int i = head[u], v; v = edges[i].to, i; i = edges[i].next)
        if (v ^ f[u][0]) {
            f[v][0] = u, dep[v] = dep[u] + 1, dfs(v);
            for (int i = 1; i <= pcolor; ++i) cnt[u][i] += cnt[v][i];
        }
    out[u] = pdfn;
}
int pre(int x, int y) {
    for (int k = LogN; ~k; --k)
        if (dep[f[y][k]] > dep[x]) y = f[y][k];
    return y;
}
}  // namespace Gragh
namespace DS {
#define belong(x) ((x) / B)
#define lp(x) ((x)*B)
#define rp(x) (std::min((x)*B + B, n) - 1)
int c[N];
LL s[SN];
void add(int x) { ++c[x], ++s[belong(x)]; }
LL sum(int l, int r) {
    if (l > r) return 0;
    int bl = belong(l), br = belong(r);
    LL ans = 0;
    if (bl == br)
        for (int i = l; i <= r; ++i) ans += c[i];
    else {
        for (int i = rp(bl); i >= l; --i) ans += c[i];
        for (int i = bl + 1; i < br; ++i) ans += s[i];
        for (int i = lp(br); i <= r; ++i) ans += c[i];
    }
    return ans;
}
}  // namespace DS
struct Query {
    int l, r, id, d;
};
std::vector<Query> que[N];
int main() {
    input(n), input(q), B = sqrt(n + 0.5);
    static int p[N], appear[N], V;
    for (int i = 1; i <= n; ++i) input(a[i]), p[i] = a[i];
    std::sort(p + 1, p + n + 1), V = std::unique(p + 1, p + n + 1) - p - 1;
    for (int i = 1; i <= n; ++i) ++appear[a[i] = std::lower_bound(p + 1, p + V + 1, a[i]) - p];
    for (int i = 1; i <= V; ++i)
        if (appear[i] >= B) color[++pcolor] = i;
    for (int i = 1, u, v; i < n; ++i) input(u), input(v), Gragh::addEdges(u, v);
    Gragh::dfs(1);
    static LL ans[M];
    int pans = 0;
    for (int i = 1, opt, x, y, root = 1, f; i <= q; ++i) {
        input(opt), input(x);
        switch (opt) {
            case 1: root = x; break;
            case 2: input(y), ++pans;
#define intree(t) (in[t] <= in[root] && in[root] <= out[t])
                for (int k = 1; k <= pcolor; ++k) {
                    LL ansx = x == root   ? cnt[1][k]
                              : intree(x) ? cnt[1][k] - cnt[Gragh::pre(x, root)][k]
                                          : cnt[x][k];
                    LL ansy = y == root   ? cnt[1][k]
                              : intree(y) ? cnt[1][k] - cnt[Gragh::pre(y, root)][k]
                                          : cnt[y][k];
                    ans[pans] += ansx * ansy;
                }
                if (x == root) {
                    if (y == root) que[n].push_back({1, n, pans, 1});
                    if (intree(y))
                        f = Gragh::pre(y, root), que[n].push_back({1, in[f] - 1, pans, 1}),
                        que[n].push_back({out[f] + 1, n, pans, 1});
                    else
                        que[n].push_back({in[y], out[y], pans, 1});
                } else if (intree(x)) {
                    const int g = Gragh::pre(x, root);
                    if (y == root)
                        que[in[g] - 1].push_back({1, n, pans, 1}), que[out[g]].push_back({1, n, pans, -1}),
                            que[n].push_back({1, n, pans, 1});
                    else if (intree(y))
                        f = Gragh::pre(y, root), que[in[g] - 1].push_back({1, in[f] - 1, pans, 1}),
                        que[in[g] - 1].push_back({out[f] + 1, n, pans, 1}),
                        que[out[g]].push_back({1, in[f] - 1, pans, -1}),
                        que[out[g]].push_back({out[f] + 1, n, pans, -1}),
                        que[n].push_back({1, in[f] - 1, pans, 1}), que[n].push_back({out[f] + 1, n, pans, 1});
                    else
                        que[in[g] - 1].push_back({in[y], out[y], pans, 1}),
                            que[out[g]].push_back({in[y], out[y], pans, -1}),
                            que[n].push_back({in[y], out[y], pans, 1});
                } else {
                    if (y == root)
                        que[in[x] - 1].push_back({1, n, pans, -1}), que[out[x]].push_back({1, n, pans, 1});
                    else if (intree(y))
                        f = Gragh::pre(y, root), que[in[x] - 1].push_back({1, in[f] - 1, pans, -1}),
                        que[in[x] - 1].push_back({out[f] + 1, n, pans, -1}),
                        que[out[x]].push_back({1, in[f] - 1, pans, 1}),
                        que[out[x]].push_back({out[f] + 1, n, pans, 1});
                    else
                        que[in[x] - 1].push_back({in[y], out[y], pans, -1}),
                            que[out[x]].push_back({in[y], out[y], pans, 1});
                }
                break;
        }
    }
    for (int i = 1; i <= n; ++i) {
        for (int u : upd[i]) DS::add(u);
        for (Query& u : que[i]) ans[u.id] += u.d * DS::sum(u.l, u.r);
    }
    for (int i = 1; i <= pans; ++i) printf("%lld\n", ans[i]);
    return 0;
}