【JZOJ6239】【20190629】智慧树

题目

一颗$n$个节点的树，每个点有一个权值$a_i$

询问树上连通块权值之和对 $m$ 取模为$ x $ 的方案数

答案对$950009857$ 取模，满足$m | 950009856$

空间$32 \ M$

题解

考虑$F_i(x)$为i为根的连通块的生成函数，$G(x)$为答案的生成函数

\[\begin{align}
F_i(x) \ = \ (\prod F_{son(i)}+1)x^{a_i} \\
G(x) \ = \ \sum F_i(x) \\
\end{align}
\]
一个比较naive的想法就是每次都dft,idft之后手动将系数取模，但其实没有必要
可以直接对每个点求出dft，O(nm)求出点乘后的结果，最后做一次idft之后的到的答案就是循环卷积

注意每个点的形式是$x^a_i$，预处理单位根的次幂可以直接得到dft
对于空间，重链剖分之后把一个点的数组直接传给重儿子，一个点dfs完之后回收一下

最大引用数组就是一条链上的轻链个数
时间复杂度:$O(nm + m \ log \ m)$ ; 空间复杂度:$O(m \ log \ n)$
为什么idft之后的结果是循环卷积呢？

\[\begin{align}
&考虑一个次数界为n的式子\{A_i\}在m次下的点值表达：\{dft(A)_i\}\\
&设\{B_i\}为idft之后的结果\\
&B_r = \sum_{j=0}^{m-1}\frac{1}{m}dft(A)_j \omega_m^{-jr}\\
&= \sum_{j=0}^{m-1}\sum_{k=0}^{n-1}\frac{1}{m}A_k \omega_m^{j(k-r)}\\
&= \sum_{k=0}^{n-1}A_k (\sum_{j=0}^{m-1}\frac{1}{m}\omega_m^{j(k-r)})\\
&运用求和引理，B_r = \sum_{j=0}^{n-1}A_k[m|j-r]\\
&即m次意义下的循环卷积
\end{align}
\]

Code

#include<bits/stdc++.h>

#define mod 950009857

using namespace std;

const int N=1<<18,M=14,R=7;

int n,m,a[N],o=1,hd[N],sz[N],sn[N],f[M][N],q[M],head,tail,now,id[N],g[N];

int t1[N],t2[N],len,L,rev[N],w[N],W[N];

struct Edge{int v,nt;}E[N<<1];

void adde(int u,int v){

	E[o]=(Edge){v,hd[u]};hd[u]=o++;

	E[o]=(Edge){u,hd[v]};hd[v]=o++;

}

void inc(int&x,int y){x+=y;if(x>=mod)x-=mod;}

void dfsA(int u,int fa){

	sz[u]=1;

	for(int i=hd[u];i;i=E[i].nt){

		int v=E[i].v;

		if(v==fa)continue;

		dfsA(v,u);

		sz[u]+=sz[v];

		if(sz[v]>sz[sn[u]])sn[u]=v;

	}

}

int get(){

	if(head==tail)q[tail++]=++now;

	if(tail==21)tail=0;

	int re=q[head++],*F=f[re];

	for(int i=0;i<m;++i)F[i]=1;

	if(head==21)head=0;

	return re;

}

void pop(int x){

	q[tail++]=x;

	if(tail==21)tail=0;

}

void dfsB(int u,int fa){

	if(sn[u]){

		dfsB(sn[u],u);

		id[u]=id[sn[u]];

		int *F=f[id[u]];

		for(int i=0;i<m;++i)inc(F[i],1);

	}else id[u]=get();

	int *F=f[id[u]];

	for(int i=0,t=0;i<m;++i){

		F[i]=1ll*F[i]*w[t]%mod;

		t+=a[u];if(t>=m)t-=m;

	}

	for(int i=hd[u];i;i=E[i].nt){

		int v=E[i].v;

		if(v==fa||v==sn[u])continue;

		dfsB(v,u);

		int *G=f[id[v]];

		for(int j=0;j<m;++j)F[j]=1ll*F[j]*(G[j]+1)%mod;

		pop(id[v]);

	}

	for(int i=0;i<m;++i)inc(g[i],F[i]);

}

int pw(int x,int y){

	int re=1;if(y<0)y+=mod-1;

	for(;y;y>>=1,x=1ll*x*x%mod)

		if(y&1)re=1ll*re*x%mod;

	return re;

}

void ntt(int*A,int F){

	for(int i=0;i<len;++i)if(i<rev[i])swap(A[i],A[rev[i]]);

	for(int i=1;i<len;i<<=1){

		int wn=pw(R,F*(mod-1)/i/2);

		for(int j=0;j<len;j+=i<<1){

			int t=1;

			for(int k=0;k<i;++k,t=1ll*t*wn%mod){

				int x=A[j+k],y=1ll*t*A[j+k+i]%mod;

				A[j+k]=(x+y)%mod;A[j+k+i]=(x-y+mod)%mod;

			}

		}

	}

	if(!~F){

		int iv=pw(len,mod-2);

		for(int i=0;i<len;++i)A[i]=1ll*A[i]*iv%mod;

	}

}

int main(){

	freopen("tree.in","r",stdin);

	freopen("tree.out","w",stdout);

	scanf("%d%d",&n,&m);

	w[0]=1;w[1]=pw(R,(mod-1)/m);

	for(int i=2;i<=m;++i)w[i]=1ll*w[i-1]*w[1]%mod;

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

		adde(u,v);

	}

	dfsA(1,0);

	dfsB(1,0);

	for(len=1,L=0;len<=(m-1)*3;len<<=1,++L);

	for(int i=1;i<len;++i)rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));

	for(int i=0;i<=(m-1)<<1;++i){

		W[i]=w[1ll*i*(i-1)/2%m];

		t2[i]=pw(W[i],mod-2);

	}

	for(int i=0;i<m;++i)t1[i]=1ll*g[m-1-i]*W[m-1-i]%mod;

	ntt(t1,1);ntt(t2,1);

	for(int i=0;i<len;++i)t1[i]=1ll*t1[i]*t2[i]%mod;

	ntt(t1,-1);

	int iv=pw(m,mod-2);

	for(int i=0;i<m;++i)printf("%d ",1ll*iv*W[i]%mod*t1[m-1+i]%mod);

	return 0;

}

巴特西

【JZOJ6239】【20190629】智慧树

题目

题解

Code

最新文章

热门文章