10.15 NOTE

YCYZ 5015: 2.树上求和(tree)

题目传送门

卡了我一万年。。。。。。

思路

显然可以得到一个朴素算法：一遍 DFS 预处理树上倍增信息与每个结点到根的答案，求值时枚举点对，根据 LCA 求得答案。
时间复杂度 \(O(n^2log n)\)，显然无法通过本题。

关键观察：转换思路，题目等价于，对于每一个颜色，有多少个点对的路径上包含这个颜色？
这相当于：考虑到某个颜色，删去这个颜色的所有结点，剩下的每个连通块内点对数量求和就是不包含这个颜色的点对，答案用总点对数量减去这个值即可。

那我们就可以开始树形 DP 了。。。时间复杂度 \(O(n)\)，显然可以通过本题。

给出柿子：

\[\large ans=\sum_{c}(C_{n}^{2}-\sum_{sz(\textup{\textrm{连通块大小}})}C_{sz_{c,i}}^2) \]

要注意的细节问题

由于有效答案是每个结点的每一个子节点的子树 \(-\) 该子树中被同一种颜色结点覆盖的大小，因此不可以直接对这个结点的大小动刀。

那么我们可以维护一个数组 \(sum_{c_i}\) 来表示目前对于颜色 \(c\) 已经处理完的子树大小，在每一次更新完儿子的连通块大小后用 \(pre\) 记录当前已经处理完的子树大小，从而方便地得到每一个儿子所对应的连通块大小。
就这个玩意卡了我半天。。。。

总结

转换思路/观察同一问题的其他表达方式

最主要的其实还是正难则反，但是，有的时候，当一个题目给出的显然元素的贡献难以求得时，不妨试试考虑另一种关键元素的贡献。

Code

#include<bits/stdc++.h>
#define Iseri namespace
#define Nina std
#define Kawaragi int
#define Momoka main
#define ls(p) (p<<1)
#define rs(p) (p<<1|1)
#define ll long long
#define ull unsigned long long
#define endl "\n"
#define pii pair<ll,ll>
const int maxn=200005;
const int inf=0x3f3f3f3f;

using Iseri Nina;

inline ll read(){
	ll x=0,f=1;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-48;ch=getchar();}
	return x*f;
}

//===========================================================

ll n,c[maxn],sz[maxn],x,y,ans,cnt,sum[maxn],vis[maxn];
vector<ll>v[maxn];

inline void dfs(ll u,ll fa){
	sz[u]=1;
	ll pre=sum[c[u]],cl=0,szv;

	for(auto i:v[u]){
		if(i==fa)continue;
		dfs(i,u);

		sz[u]+=sz[i];
		ll tmp=sum[c[u]]-pre;//表示在目前这个儿子中已被覆盖的结点数量
		szv=sz[i]-tmp;//表示在目前这个儿子中未被同一颜色阻断的连通块大小
		pre=sum[c[u]];//更新已经处理完的覆盖大小
		ans-=(szv*(szv-1))/2;//更新答案
		cl+=tmp;//防止重复计算sz[u]
	}
	sum[c[u]]+=sz[u]-cl;//把sz[u]存入sum
}

Kawaragi Momoka(){
	n=read();
	for(ll i=1;i<=n;i++){
		c[i]=read();
		if(!vis[c[i]])vis[c[i]]=1,cnt++;
		sz[i]=1;
	}
	for(ll i=1;i<n;i++){
		x=read(),y=read();
		v[x].push_back(y);
		v[y].push_back(x);
	}

	ans=(cnt*n*(n-1))/2;
	dfs(1,0);

	for(ll i=0;i<=n;i++){
		if(vis[i]==0)continue;
		ll tmp=n-sum[i];
		ans-=(tmp*(tmp-1))/2;
	}

	printf("%lld\n",ans);
	return 0;
}

这是绝对的好题……