题目
给你一棵树,每个点有点权。
一个点的贡献为:所有以它为编号最小的(即如果有两个重心,它是编号小那个)重心的大小为 k k k的树连通块的方案数乘点权。
问贡献和。
n ≤ 5 ∗ 1 0 4 n\le 5*10^4 n≤5∗104
k ≤ 500 k\le 500 k≤500
首先 O ( n k 2 ) O(nk^2) O(nk2)的没有人想不到吧,直接树上背包(这里时间复杂度 O ( n k ) O(nk) O(nk))+换根(这里时间复杂度 O ( n k 2 ) O(nk^2) O(nk2))。
这个方法太low,原因是我们分别对于每个点求出以它为重心的方案数。
于是有了下面这个DP:设 g i , j g_{i,j} gi,j表示以 i i i为根的子树,选了大小为 j j j的连通块,此时备选重心的贡献和。
设 f i , j f_{i,j} fi,j表示以 i i i为根的子树,选了大小为 j j j的连通块的方案数。
f f f的转移显然。现在看 g g g。
对于遍历到的某个点 i i i,假如 i i i不作为重心,那就直接转移 g g g。即 g i ← g i f j + g j f i g_i\leftarrow g_if_j+g_jf_i gi←gifj+gjfi。
假如 i i i作为重心,那么一定要满足:选出的连通块中, i i i的儿子的子树的大小不超过 k 2 \frac{k}{2} 2k,并且 k − x 子 树 大 小 k-x子树大小 k−x子树大小也不超过 k 2 \frac{k}{2} 2k。我们先不用管 i i i子树的补集,如果满足这个条件,那么 i i i就可以成为一个备选重心。
具体来说(为了方便用生成函数表示啦),算出 a i x ∏ j ∈ s o n ( i ) ( f j m o d x k 2 + 1 ) a_ix\prod_{j\in son(i)} (f_{j}\mod x^{\frac{k}{2}+1}) aix∏j∈son(i)(fjmodx2k+1),然后保留满足 k − t ≤ k 2 k-t\le \frac{k}{2} k−t≤2k的 x t x^t xt项。
两种情况加起来就可以完成对 g g g的转移啦。
至于 2 ∣ k 2|k 2∣k时可能会出现的有两个重心的情况,乱搞一下将不该加上的那个贡献减去即可。
时间复杂度为 O ( n k ) O(nk) O(nk)。对于它的分析可以归结成树上大小为 k k 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;
}