題目傳送門

題目大意

給出一個 \(n\) 個點的樹，每個點有顏色，定義 \(\text{dis}(u,v)\) 為兩個點之間不同顏色個數，有 \(m\) 次修改，每次將某個點的顏色進行更改，在每次操作後求出：

\[\sum_{i=1}^{n}\sum_{j=1}^{n}\text{dis}(i,j) \]

\(n,m\le 4\times 10^5\)

思路

好妙的一道題啊！！！看 \(\text{yyb}\) 神仙的部落格看到的，花了我一個晚上。。。而且還是看題解看懂的。。。

首先我們可以想到，肯定是對於每一種顏色進行考慮，但是考慮出現的方案數顯然不好搞，於是我們容斥一下就變成了總方案數減去沒有出現過的方案數。然後我們發現如果我們把當前顏色設為白色，不同顏色設為黑色，那麼答案就是黑色連通塊大小平方之和。於是，問題就是如何求這個。

我們有一個人盡皆知的 \(\text{trick}\) ，就是說我們可以黑點往父節點連邊，然後實際聯通塊就是樹上連通塊除去根了。然後這個東西就可以用 \(\text{LCT}\) 進行維護了，只需要維護虛子樹資訊即可。

時間複雜度 \(\Theta((n+m)\log n)\) ，具體實現見程式碼。

\(\texttt{Code}\)

#include <bits/stdc++.h>
using namespace std;

#define PII pair<int,int>
#define Int register int
#define ll long long
#define MAXN 400005

template <typename T> inline void read (T &t){t = 0;char c = getchar();int f = 1;while (c < '0' || c > '9'){if (c == '-') f = -f;c = getchar();}while (c >= '0' && c <= '9'){t = (t << 3) + (t << 1) + c - '0';c = getchar();} t *= f;}
template <typename T,typename ... Args> inline void read (T &t,Args&... args){read (t);read (args...);}
template <typename T> inline void write (T x){if (x < 0){x = -x;putchar ('-');}if (x > 9) write (x / 10);putchar (x % 10 + '0');}

ll Num;

struct LCT{
#define ls(x) son[x][0]
#define rs(x) son[x][1]
	int fa[MAXN],siz[MAXN],vsz[MAXN],son[MAXN][2];ll ssz[MAXN];
	ll val (int x){return 1ll * siz[x] * siz[x];}
	bool rnk (int x){return son[fa[x]][1] == x;}
	bool Isroot (int x){return son[fa[x]][0] != x && son[fa[x]][1] != x;}
	void Pushup (int x){siz[x] = siz[ls (x)] + siz[rs (x)] + vsz[x] + 1;}
	void rotate (int x){
		int y = fa[x],z = fa[y],k = rnk (x),w = son[x][!k];
		if (!Isroot (y)) son[z][rnk (y)] = x;son[x][!k] = y,son[y][k] = w;
		if (w) fa[w] = y;fa[x] = z,fa[y] = x;
		Pushup (y),Pushup (x);
	}
	void Splay (int x){
		while (!Isroot (x)){
			int y = fa[x];
			if (!Isroot (y)) rotate (rnk (x) == rnk (y) ? y : x);
			rotate (x);
		}
		Pushup (x);
	}
	void Access (int x){
		for (Int y = 0;x;x = fa[y = x])
			Splay (x),vsz[x] += siz[rs (x)],vsz[x] -= siz[y],ssz[x] += val (rs (x)),ssz[x] -= val (y),rs(x) = y,Pushup (x);
	}
	int findroot (int x){Access (x),Splay (x);while (ls (x)) x = ls (x);Splay (x);return x;}
	void link (int x,int y){
		Access (x),Num -= ssz[x];
		int z = findroot (y);Splay (z),Num -= val (rs (z));
		fa[x] = y,Splay (y),vsz[y] += siz[x],ssz[y] += val (x);
		Pushup (y),Access (x),Splay (z),Num += val (rs (z));
	}
	void cut (int x,int y){
		Access (x),Num += ssz[x];
		int z = findroot (y);Access (x),Splay (z),Num -= val (rs (z));
		Splay (x),son[x][0] = fa[son[x][0]] = 0;
		Pushup (x),Splay (z),Num += val (rs (z));
	}
}Tree;

vector <int> E[MAXN];
vector <PII> V[MAXN];
int n,m,c[MAXN],fa[MAXN],col[MAXN];ll ans[MAXN];

void dfs (int u,int par){
	fa[u] = par;
	for (Int v : E[u]) if (v ^ par) dfs (v,u);
}

signed main(){
	read (n,m);
	for (Int i = 1;i <= n;++ i) read (c[i]);
	for (Int i = 2,u,v;i <= n;++ i) read (u,v),E[u].push_back (v),E[v].push_back (u);
	for (Int i = 1;i <= n;++ i) V[c[i]].push_back (make_pair (0,i));
	for (Int i = 1,u,v;i <= m;++ i) read (u,v),V[c[u]].push_back (make_pair (i,u)),V[c[u] = v].push_back (make_pair (i,u));
	dfs (1,n + 1);
	for (Int i = 1;i <= n + 1;++ i) Tree.Pushup (i);
	for (Int i = 1;i <= n;++ i) Tree.link (i,fa[i]);
	for (Int i = 1;i <= n;++ i){
		ll lst = 0;
		for (auto v : V[i]){
			int t = v.first,u = v.second;
			if (col[u]) Tree.link (u,fa[u]);else Tree.cut (u,fa[u]);
			col[u] ^= 1,ans[t] += 1ll * n * n - Num - lst,lst = 1ll * n * n - Num;
		}
		for (auto v : V[i]){
			int u = v.second;
			if (col[u]) col[u] ^= 1,Tree.link (u,fa[u]);
		}
	}  
	for (Int i = 1;i <= m;++ i) ans[i] += ans[i - 1];
	for (Int i = 0;i <= m;++ i) write (ans[i]),putchar ('\n');
 	return 0;
}