codeforces1172 E. Nauuo and ODT LCT
分析
題目大意：
給一棵每個節點有顏色的樹，求樹上所有路徑不同顏色個數和。帶修改。
一個顯然的思路是分顏色考慮，對於每種顏色，考慮沒有貢獻的路徑，那麼就是去掉這些顏色的點，剩下的每個連通塊內部的點對。也就是連通塊的大小平方和條路徑沒有貢獻。。把這些顏色的點全部看成白色，其餘看成黑色，問題轉化成給你初始是黑色的樹，要求維護：1.單點顏色取反2.求黑點連通塊大小平方和。
有一種神奇的操作，先轉有根樹，是每個黑點與其父親連邊，並且新建一個虛的白點，讓根向它連邊。這樣連邊的話，每個連通塊的根一定是白點，那麼答案就是所有白點的兒子子樹大小平方和。
取反操作相當於是斷邊和連邊。
上LCT即可。
代碼
#include<bits/stdc++.h>
#define mp std::make_pair
#define fi first
#define se second
const int N = 4e5 + 10;
int ri() {
char c = getchar(); int x = 0, f = 1; for(;c < '0' || c > '9'; c = getchar()) if(c == '-') f = -1;
for(;c >= '0' && c <= '9'; c = getchar()) x = (x << 1) + (x << 3) - '0' + c; return x * f;
}
long long ans[N];
int c[N], pr[N], nx[N << 1], to[N << 1], tp, n, m;
void add(int u, int v) {to[++tp] = v; nx[tp] = pr[u]; pr[u] = tp;}
void adds(int u, int v) {add(u, v); add(v, u);}
std::vector<int> ini[N];
std::vector<std::pair<int, int> > mod[N];
namespace LCT {
long long ans = 0, isz2[N];
int isz[N], sz[N], rfa[N], fa[N], ch[N][2];
int tp, st[N << 1]; bool wt[N];
#define ls ch[p][0]
#define rs ch[p][1]
bool wh(int p) {return ch[fa[p]][1] == p;}
bool Ir(int p) {return ch[fa[p]][0] != p && ch[fa[p]][1] != p;}
long long sqr(long long n) {return n * n;}
void Build(int u, int f) {
fa[u] = f; rfa[u] = f; sz[u] = 1; isz2[u] = 0;
for(int i = pr[u]; i;i = nx[i])
if(to[i] != f)
Build(to[i], u), sz[u] += sz[to[i]], isz2[u] += sqr(sz[to[i]]);
isz[u] = sz[u] - 1;
}
void Up(int p) {sz[p] = isz[p] + sz[ls] + sz[rs] + 1;}
void Rotate(int p) {
int f = fa[p], g = fa[f], c = wh(p);
if(!Ir(f)) ch[g][wh(f)] = p; fa[p] = g;
ch[f][c] = ch[p][c ^ 1]; if(ch[f][c]) fa[ch[f][c]] = f;
ch[p][c ^ 1] = f; fa[f] = p; Up(f);
}
void Splay(int p) {
for(;!Ir(p); Rotate(p))
if(!Ir(fa[p]))
Rotate(wh(fa[p]) == wh(p) ? fa[p] : p);
Up(p);
}
void Access(int u) {
for(int pr = 0, p = u; p;pr = p, p = fa[p]) {
Splay(p);
if(rs) {
isz[p] += sz[rs];
isz2[p] += sqr(sz[rs]);
}
if(pr) {
isz[p] -= sz[pr];
isz2[p] -= sqr(sz[pr]);
}
rs = pr;
Up(p);
}
Splay(u);
}
int Find(int p) {
Access(p);
for(;ls; p = ls) ;
return p;
}
long long Qry(int u) {
int p = Find(u);
return Access(p), isz2[p];
}
void Cut(int u) {
ans -= Qry(u);
Access(u); int v = ch[u][0];
ch[u][0] = fa[v] = 0; Up(u);
ans += Qry(u) + Qry(v);
}
void Link(int u, int v) {
ans -= Qry(u) + Qry(v);
Access(u); Access(v);
fa[u] = v;
isz[v] += sz[u];
isz2[v] += sqr(sz[u]);
Up(v);
ans += Qry(u);
}
void Rev(int u, bool p) {
if(p)
st[++tp] = u;
if(!wt[u]) {
wt[u] = true;
Cut(u);
}
else {
wt[u] = false;
Link(u, rfa[u]);
}
}
void Restore() {
for(int i = 1;i <= tp; ++i)
if(wt[st[i]])
Rev(st[i], 0);
tp = 0;
}
}
int main() {
n = ri(); m = ri();
for(int i = 1;i <= n; ++i)
c[i] = ri(), ini[c[i]].push_back(i);
for(int i = 1;i < n; ++i)
adds(ri(), ri());
adds(1, n + 1);
LCT::Build(n + 1, 0); LCT::ans = 1LL * n * n;
long long all = LCT::ans;
for(int i = 1;i <= m; ++i) {
int u = ri(), x = ri();
if(c[u] == x)
continue;
mod[c[u]].push_back(mp(u, i));
c[u] = x;
mod[c[u]].push_back(mp(u, i));
}
for(int x = 1;x <= n; ++x) {
for(auto i : ini[x])
LCT::Rev(i, 1);
ans[0] += all - LCT::ans;
for(auto i : mod[x]) {
ans[i.se] -= all - LCT::ans;
LCT::Rev(i.fi, 1);
ans[i.se] += all - LCT::ans;
}
LCT::Restore();
}
for(int i = 1;i <= m; ++i)
ans[i] += ans[i - 1];
for(int i = 0;i <= m; ++i)
printf("%lld\n", ans[i]);
return 0;
}