牛客练习赛81 D.小 Q 与树(点分治+容斥)

Posted 2021-06-25 issue是fw

LINK

$n$个点的树,点$i$的点权是$a_i$

求$$\sum _{u=1}^n\sum _{v=1}^n\mathrm{m}\mathrm{i}\mathrm{n}({\mathrm{a}}_{\mathrm{u}},{\mathrm{a}}_{\mathrm{v}})\ast \mathrm{d}\mathrm{i}\mathrm{s}(\mathrm{u},\mathrm{v})$$

---

树上路径问题,考虑点分治

找出重心$u$,然后计算所有通过点$u$的路径贡献

然后从原图中消去点$u$,然后对各个连通块进行同样的分治.

通过点$u$的路径贡献怎么算…

先爆搜出所有从$u$出发的路径,存一个$pair$

第一关键字是这条路径的非$u$端点点权,第二关键字是这条路径的长度

那么把这些路径排个序,通过点$u$的路径无非两种

$Ⅰ.$只从$u$出发的路径

$Ⅱ.$由两条从$u$出发的路径拼接而成

$Ⅰ$类路径可以在爆搜的时候顺便算掉,$Ⅱ$类路径排序之后也很好算,第$i$条路径的点权比$[i+1,top]$都小,所以直接去掉了式子中的$\min$,暴力算就好了

但是这样有重复,比如$v_1,v_2$都来自$u$的同一颗子树$v$内,然而在上述的计算方式中却把他们拼接了

所以此时应该减去子树的贡献,这就是一个容斥的过程

值得一提的是,计算子树时的贡献初始路径长度为$1$,因为我们现在在容斥,应该看作所有路径从$u$出发

#include <bits/stdc++.h>
using namespace std;
const int maxn = 4e5+10;
const int mod = 998244353;
int n,a[maxn],ans;
vector<int>vec[maxn];
int siz[maxn],vis[maxn],mx[maxn],root,sumn;
typedef pair<int,int>p;
p res[maxn]; int top;
void getroot(int u,int fa)
{
	siz[u] = 1, mx[u] = 0;
	for( auto v:vec[u] )
	{
		if( vis[v] || v==fa )	continue;
		getroot(v,u);
		siz[u] += siz[v];
		mx[u] = max( mx[u],siz[v] );
	}
	mx[u] = max( mx[u],sumn-siz[u] );
	if( mx[u]<mx[root] )	root = u;
}
void dfs(int u,int fa,int dep)
{
	res[++top] = p( a[u],dep );
	for(auto v:vec[u] )
	{
		if( v==fa || vis[v] )	continue;
	//	res[++top] = p( a[v],dep+1 );
		dfs( v,u,dep+1 );
	}
}
//上下两行注释这么写也是对的. 
int calc(int u,int initdep)
{
	long long ans = 0, suf = 0;
	top = 0; dfs( u,u,initdep );
	sort( res+1,res+1+top );
	for(int i=top;i>=1;i--)
	{
//		ans = ( ans+min( res[i].first,a[u] )*( initdep+res[i].second ) )%mod;
		ans = ( ans+1ll*res[i].first*(suf+1ll*res[i].second*(top-i)%mod)%mod )%mod;
		suf = ( suf+res[i].second )%mod;
	}
	return ans;
}
void solve(int u)
{
	vis[u] = 1; ans = ( ans+calc(u,0) )%mod;
	for(auto v:vec[u] )
	{
		if( vis[v] )	continue;
		ans = ( ans-calc(v,1) )%mod;
		sumn = siz[v], mx[root=0] = n+1;
		getroot(v,0); solve( root );
	}
}
int main()
{
	scanf("%d",&n );
	for(int i=1;i<=n;i++)	scanf("%d",&a[i] );
	for(int i=1;i<n;i++)
	{
		int l,r; scanf("%d%d",&l,&r);
		vec[l].push_back( r ); vec[r].push_back( l );
	}
	sumn = n, mx[root=0] = n+1;
	getroot(1,0); solve(1);
	printf("%d",(1ll*ans*2%mod+mod)%mod );
}