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;
}