题解 CF1172E Nauuo and ODT

题目传送门

题目大意

给出一个 \(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;
}