一顆\(n\)個節點的樹,每一個點有一個權值\(a_i\)c++
詢問樹上連通塊權值之和對 \(m\) 取模爲$ x $ 的方案數數組
答案對\(950009857\) 取模,知足\(m | 950009856\)函數
空間\(32 \ M\)spa
考慮\(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} \]code
一個比較naive的想法就是每次都dft,idft以後手動將係數取模,但其實沒有必要get
能夠直接對每一個點求出dft,O(nm)求出點乘後的結果,最後作一次idft以後的到的答案就是循環卷積it
注意每一個點的形式是\(x^a_i\),預處理單位根的次冪能夠直接獲得dftclass
對於空間,重鏈剖分以後把一個點的數組直接傳給重兒子,一個點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} \]
#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; }