CF960E Alternating Tree

一、题目

点此看题

二、解法

一开始想对于每个根统计其子树贡献，但是难。换一个角度，由于加法是相对独立的，于是考虑每个点的贡献。

有向路径有点麻烦，我们来分类讨论：

• $x$子树内-> $x$-> $x$子树外（包括 $x$）
• $x$子树内（包括 $x$）-> $x$-> $x$的另一个子树
• $x$子树外-> $x$-> $x$的子树（包括 $x$）

第我们需要处理出：子树内的和它距离为偶数的点数 $fd$，子树内和它距离为奇数的点数 $gd$，子树外的和它距离为偶数的点数 $fu$，子树内和外距离为奇数的点数 $gu$，子树内的点数，对应上面的情况计算：

• $a[u]\times(fd[u]-gd[u])\times(n-siz[u]+1)$
• $a[u]\times(gd[v]-fd[v])\times(siz[u]-siz[v]-1)+a[u]\times(siz[u]-1)$
• $a[u]\times (fu[u]-gu[u])\times siz[u]$

然后跑两遍 $dfs$就可以辣

#include <cstdio>
#define int long long
const int M = 200005;
const int MOD = 1e9+7;
int read()
{
    int x=0,flag=1;char c;
    while((c=getchar())<'0' || c>'9') if(c=='-') flag=-1;
    while(c>='0' && c<='9') x=(x<<3)+(x<<1)+(c^48),c=getchar();
    return x*flag;
}
int n,tot,f[M],a[M],siz[M];
int ans,fd[M],gd[M],fu[M],gu[M];
struct edge
{
	int v,next;
}e[2*M];
void dfs(int u,int fa)
{
	siz[u]=fd[u]=1;
	for(int i=f[u];i;i=e[i].next)
	{
		int v=e[i].v;
		if(v==fa) continue;
		dfs(v,u);
		siz[u]+=siz[v];
		fd[u]+=gd[v];
		gd[u]+=fd[v];
	}
	ans=(ans+a[u]*(fd[u]-gd[u])%MOD*(n-siz[u]+1))%MOD;
	for(int i=f[u];i;i=e[i].next)
	{
		int v=e[i].v;
		if(v^fa)
			ans=(ans+a[u]*(gd[v]-fd[v])%MOD*(siz[u]-siz[v]-1)%MOD)%MOD;
	}
	ans=(ans+a[u]*(siz[u]-1))%MOD;
}
void dfs2(int u,int fa)
{
	if(fa)
	{
		fu[u]=gu[fa]+gd[fa]-fd[u];
		gu[u]=fu[fa]+fd[fa]-gd[u];
	}
	for(int i=f[u];i;i=e[i].next)
	{
		int v=e[i].v;
		if(v==fa) continue;
		dfs2(v,u);
	}
	ans=(ans+a[u]*(fu[u]-gu[u])%MOD*siz[u]%MOD)%MOD;
}
signed main()
{
	n=read();
	for(int i=1;i<=n;i++)
		a[i]=read();
	for(int i=1;i<n;i++)
	{
		int u=read(),v=read();
		e[++tot]=edge{v,f[u]},f[u]=tot;
		e[++tot]=edge{u,f[v]},f[v]=tot;
	}
	dfs(1,0);
	dfs2(1,0);
	printf("%lld\n",(ans+MOD)%MOD);
}