牛客练习赛60 E 旗鼓相当的对手

LINK：旗鼓相当的对手

考场上遇到这种简单的树形dp优化题目我真的不一定能写出来..

虽然这道题思路很简单 设f[i][j]表示距i距离为j的点的个数 g[i][j]表示距i距离为j的点权和。

可以发现我们转枚举子树转移的时候可以得到某个点的答案。其实这道题让我们求的是以x为根的所有子树之间的答案。

这两个数组转移以深度转移 所以复杂度为n^2 长链刨分一下即可O(n)。

但是我并不会长链刨分时的指针写法。所以考虑其他做法。

可以发现如果维护这两个数组的转移必然n^2 这个时候我们考虑直接给全局贡献 即每个点x以d[x]为深度的贡献。

这样我们就不需要维护距某个点的深度了 直接查d[x]+k的距离即可。

但是这需要子树内外的差分。暴力差分还是n^2.

考虑dsu on tree 暴力把轻儿子删掉之后统计重儿子 然后重儿子的代价保留。

对于某个点我们再次遍历所有轻儿子一边统计答案一边累计数组。

最后如果这个点也是轻儿子那么再次情况即可。

一个点距根节点只有logn条轻边 所以这样做复杂度nlogn.

const int MAXN=100010;
int n,len,k;ll ans[MAXN],w[MAXN],sum;
int son[MAXN],a[MAXN],d[MAXN],f[MAXN],cnt[MAXN],sz[MAXN];
int lin[MAXN],ver[MAXN<<1],nex[MAXN<<1];
inline void add(int x,int y)
{
	ver[++len]=y;nex[len]=lin[x];lin[x]=len;
	ver[++len]=x;nex[len]=lin[y];lin[y]=len;
}
inline void dfs(int x,int father)
{
	sz[x]=1;f[x]=father;
	d[x]=d[father]+1;
	go(x)
	{
		if(tn==father)continue;
		dfs(tn,x);
		sz[x]+=sz[tn];
		if(sz[son[x]]<sz[tn])son[x]=tn;
	}
}
inline void update(int x,int fa,int op)
{
	int ww=k+2*d[fa]-d[x];
	if(ww>0&&op>0&&x!=fa)
	{
		sum=sum+(ll)cnt[ww]*a[x];
		sum=sum+w[ww];
	}
	go(x)if(tn!=f[x])update(tn,fa,op);
}
inline void add(int x,int fa,int op)
{
	cnt[d[x]]+=op;w[d[x]]+=op*a[x];
	go(x)if(tn!=f[x])add(tn,fa,op);
}
inline void dp(int x,int op)
{
	go(x)if(tn!=f[x]&&tn!=son[x])dp(tn,0);//先处理轻儿子.
	if(son[x])dp(son[x],1);//处理重儿子.
	go(x)if(tn!=f[x]&&tn!=son[x])update(tn,x,1),add(tn,x,1);
	++cnt[d[x]];w[d[x]]+=a[x];ans[x]=sum;sum=0;
	if(!op)add(x,x,-1);
}
int main()
{
	freopen("1.in","r",stdin);
	get(n);get(k);
	rep(1,n,i)get(a[i]);
	rep(1,n-1,i)add(read(),read());
	dfs(1,0);dp(1,1);
	rep(1,n,i)printf("%lld ",ans[i]);
	return 0;
}