JZOJ6815. 【2020.10.06提高组模拟】树的重心

Description

• 求大小为$n$的树上所有大小为$k$的连通块的重心权值之和（如果有两个重心，取编号小的那个）。

• $n\le 5e4,k\le 500$

Solution

• 显然可以得到一种$O(n{k}^2)$的换根DP，直接枚举子树内选择的点数转移即可。

• 经典套路：如果没有换根DP的部分，直接做大小为$k$的树形背包的复杂度实际上是$O(nk)$的（而不是$O(n{k}^2)$）。

• 证明： 首先合并两个子树$x$和$y$的时候，时间复杂度是$min(sz[x],k)\ast min(sz[y],k)$的。

  • $sz[x],sz[y]\ge k$时，考虑所有极小的子树使得$sz[i]\ge k$，那么是缩点后原树上的叶子，最多$n/k$个这样的叶子，它们之间合并$n/k$次，每一次$k^2$
  • $sz[x],sz[y]<k$时，考虑所有极大的子树使得$sz[i]<k$，那么这类合并都在这种子树里面，$sz[i]<k$，所以内部的点两两产生贡献，每一个点最多和$k$个点产生贡献，即$nk$.
  • $sz[x]\ge k,sz[y]<k$时，考虑$y$的子树内的每个节点，产生$k$的贡献，所以$nk$。

• 既然如此，记录一个$f[x][i]$和$g[x][i]$分别表示$x$子树内连通块大小为$i$的方案数和重心权值和，做背包DP即可。

#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#define min(a,b) ((a<b)?a:b)
#define maxn 50005
#define maxm 505
#define ll long long 
#define mo 1000000007
using namespace std;

int n,m,i,j,k,a[maxn],sz[maxn];
int em,e[maxn*2],nx[maxn*2],ls[maxn];
ll f[maxn][maxm],g[maxn][maxm],ans;

void read(int &x){
	x=0; char ch=getchar();
	for(;ch<'0'||ch>'9';ch=getchar());
	for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
}

void insert(int x,int y){
	em++; e[em]=y; nx[em]=ls[x]; ls[x]=em;
	em++; e[em]=x; nx[em]=ls[y]; ls[y]=em;
}

int d[maxn],fa[maxn];
void doit(int x,int p){
	sz[x]=1,f[x][1]=1;
	for(int i=ls[x];i;i=nx[i]) if (e[i]!=p) {
		int y=e[i];
		for(int j=min(m-1,sz[x]);j>=m/2;j--) for(int k=min(min(m-j,m/2),sz[y]);k;k--)
			(g[x][j+k]+=g[x][j]*f[y][k])%=mo;
		for(int j=min(m-1,sz[y]);j>=m/2;j--) for(int k=min(min(m-j,m/2),sz[x]);k;k--)
			(g[x][j+k]+=g[y][j]*f[x][k])%=mo;
		for(int j=min(m-1,sz[x]);j;j--) for(int k=min(min(m-j,m/2-((m&1)==0)),sz[y]);k;k--)
			(f[x][j+k]+=f[x][j]*f[y][k])%=mo;
		sz[x]+=sz[y];
	}
	if (m%2==0&&p) (g[x][m/2]+=f[x][m/2]*a[min(x,p)])%=mo;
	for(int i=m/2+1;i<=m;i++) if (f[x][i]) 
		(g[x][i]+=f[x][i]*a[x])%=mo;
	ans+=g[x][m];	
}
void bfs(){
	int t=0,w=1; d[1]=1;
	while (t<w){
		int x=d[++t];
		for(int i=ls[x];i;i=nx[i]) if (e[i]!=fa[x])
			d[++w]=e[i],fa[e[i]]=x;
	}
	while (w){
		int x=d[w--];
		doit(x,fa[x]);
	}
}

int main(){
//	freopen("centroid.in","r",stdin);
//	freopen("centroid.out","w",stdout);
	read(n),read(m);
	for(i=1;i<=n;i++) read(a[i]);
	for(i=1;i<n;i++) read(j),read(k),insert(j,k);
	bfs();
	printf("%lld",ans%mo);
}