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

题目

给你一棵树，每个点有点权。

一个点的贡献为：所有以它为编号最小的（即如果有两个重心，它是编号小那个）重心的大小为$k$的树连通块的方案数乘点权。

问贡献和。

$n\le 5\ast 10^4$

$k\le 500$

---

首先$O(n{k}^2)$的没有人想不到吧，直接树上背包（这里时间复杂度$O(nk)$）+换根（这里时间复杂度$O(n{k}^2)$）。

这个方法太low，原因是我们分别对于每个点求出以它为重心的方案数。

于是有了下面这个DP：设$g_{i,j}$表示以$i$为根的子树，选了大小为$j$的连通块，此时备选重心的贡献和。

设$f_{i,j}$表示以$i$为根的子树，选了大小为$j$的连通块的方案数。

$f$的转移显然。现在看$g$。

对于遍历到的某个点$i$，假如$i$不作为重心，那就直接转移$g$。即$g_i\leftarrow g_i{f}_j+g_j{f}_i$。

假如$i$作为重心，那么一定要满足：选出的连通块中，$i$的儿子的子树的大小不超过$\frac{k}{2}$，并且$k-x子树大小$也不超过$\frac{k}{2}$。我们先不用管$i$子树的补集，如果满足这个条件，那么$i$就可以成为一个备选重心。

具体来说（为了方便用生成函数表示啦），算出$a_{i}x{\prod }_{j\in son(i)}(f_j\mathrm{m}\mathrm{o}\mathrm{d}{x}^{\frac{k}{2}+1})$，然后保留满足$k-t\le \frac{k}{2}$的$x^t$项。

两种情况加起来就可以完成对$g$的转移啦。

至于$2|k$时可能会出现的有两个重心的情况，乱搞一下将不该加上的那个贡献减去即可。

时间复杂度为$O(nk)$。对于它的分析可以归结成树上大小为$k$的背包的问题，表示不太会证，但是能用不太严谨的方式去理解。

---

using namespace std;
#include <cstdio>
#include <cstring>
#include <algorithm>
#define N 50005
#define K 505
#define ll long long
#define mo 1000000007
int n,k;
int a[N];
struct EDGE{
	int to;
	EDGE *las;
} e[N*2];
int ne;
EDGE *last[N];
struct poly{
	int n,a[K];
	void set(int v){
		n=0,a[0]=v;
	}
	void move(int mx=k>>1){
		for (int i=n;i>=0;--i)
			a[i+1]=a[i];
		a[0]=0;
		n=min(n+1,mx);
	}
	void add(int v){
		(a[0]+=v)%=mo;
	}
	void add(poly &q){
		for (int i=0;i<=min(n,q.n);++i)
			(a[i]+=q.a[i])%=mo;
		if (n<q.n){
			for (int i=n+1;i<=q.n;++i)
				a[i]=q.a[i];
			n=q.n;
		}
	}
	void print(){
		for (int i=0;i<=n;++i)
			printf("%d ",a[i]);
		printf("\n");
	}
};
void multi(poly &a,poly &b,poly &c,int mx=k>>1){
	static poly t;
	t.n=min(b.n+c.n,mx);
//	memset(t.a,0,sizeof t.a);
//	for (int i=0;i<=b.n;++i)
//		for (int j=0;j<=c.n && i+j<=t.n;++j)
//			t.a[i+j]=(t.a[i+j]+(ll)b.a[i]*c.a[j])%mo;
	for (int i=0;i<=t.n;++i){
		ll s=0;
		for (int j=max(i-c.n,0);j<=i && j<=b.n;++j)
			(s+=(ll)b.a[j]*c.a[i-j])%=mo;
		t.a[i]=s;
	}
	a=t;
}
poly f[N],g[N],t;
void dp1(int x,int fa){
	f[x].set(1),f[x].move();
	g[x].set(0);
	for (EDGE *ei=last[x];ei;ei=ei->las)
		if (ei->to!=fa){
			dp1(ei->to,x);
//			g[x]=g[x]*f[ei->to]+f[x]*g[ei->to];
			multi(t,f[x],g[ei->to],k);
			multi(g[x],g[x],f[ei->to],k);
			g[x].add(t);
			multi(f[x],f[x],f[ei->to]);
		}
	f[x].add(1);
	t.set(a[x]),t.move(k);
	for (EDGE *ei=last[x];ei;ei=ei->las)
		if (ei->to!=fa){
			int tmp=f[ei->to].a[k>>1];
			if (!(k&1) && x>ei->to)
				f[ei->to].a[k>>1]=0;
			multi(t,t,f[ei->to],k);
			f[ei->to].a[k>>1]=tmp;
		}
	for (int i=0;i<=k;++i)
		if (k-i>k>>1)
			t.a[i]=0;
	if (fa && x>fa)
		t.a[k>>1]=0;
	g[x].add(t);
}
int main(){
//	freopen("in.txt","r",stdin);
//	freopen("out.txt","w",stdout);
	freopen("centroid.in","r",stdin);
	freopen("centroid.out","w",stdout);
	scanf("%d%d",&n,&k);
	for (int i=1;i<=n;++i)
		scanf("%d",&a[i]);
	for (int i=1;i<n;++i){
		int u,v;
		scanf("%d%d",&u,&v);
		e[ne]={v,last[u]};
		last[u]=e+ne++;
		e[ne]={u,last[v]};	
		last[v]=e+ne++;
	}
	dp1(1,0);
//	for (int i=1;i<=n;++i)
//		g[i].print();
	ll ans=0;
	for (int i=1;i<=n;++i)
		if (k<=g[i].n)
			ans+=g[i].a[k];
	ans%=mo;
	printf("%lld\n",ans);
	return 0;
}