【2022 CCPC 华为云计算挑战赛】1005 带权子集和（NTT 优化dp）

题目链接：

https://acm.hdu.edu.cn/contest/problem?cid=1072&pid=1005

题意：

给定大小为 n n n的可重集 A A A，给定非负整数 k , x k,x k,x。对于 A A A的一个非空子集 S S S，设 s u m ( S ) sum(S) sum(S)为集合 S S S的元素权值之和， ∣ S ∣ |S| ∣S∣为集合 S S S的元素个数。求 ∑ S ⊆ A , S ≠ ∅ s u m ( S ) k ⋅ x ∣ S ∣ \sum_{S\subseteq A,S \neq \varnothing}sum(S)^k \cdot x^{|S|} ∑S⊆A,S=∅sum(S)k⋅x∣S∣对 998244353 998244353 998244353取模的结果。

算法：

NTT优化dp。

时间复杂度：

O ( n k l o g ( k ) ) O(nklog(k)) O(nklog(k))

解法：

为了便于理解，我们先设 d p i , j , l dp_{i,j,l} dpi,j,l表示为前 i i i个，选 j j j个作为子集，在 l l l次幂下的答案。
设 p p p为前 i − 1 i-1 i−1个数中任选 j − 1 j-1 j−1个数的和， P P P为所有情况的 p p p的集合， P = { p 0 , p 1 , p 2 , . . . } P=\{p_0,p_1,p_2,...\} P={p0,p1,p2,...}。

根据题意，可以得出前 i − 1 i-1 i−1个，选 j − 1 j-1 j−1个作为子集，在 l l l次幂下的答案 d p i − 1 , j − 1 , l = ( p 0 l + p 1 l + p 2 l + . . . ) x j − 1 dp_{i-1,j-1,l}=(p_0^l+p_1^l+p_2^l+...)x^{j-1} dpi−1,j−1,l=(p0l+p1l+p2l+...)xj−1。

设 s u m l = p 0 l + p 1 l + p 2 l + . . . sum_l=p_0^l+p_1^l+p_2^l+... suml=p0l+p1l+p2l+...，可得 d p i − 1 , j − 1 , l = s u m l x j − 1 dp_{i-1,j-1,l}=sum_l x^{j-1} dpi−1,j−1,l=sumlxj−1。

考虑状态转移，考虑选 a i a_i ai给答案带来的变化：

d p i , j , l = [ ( p 0 + a i ) l + ( p 1 + a i ) l + ( p 2 + a i ) l + . . . ] x j = [ ( l 0 ) a i l s u m 0 + ( l 1 ) a i l − 1 s u m 1 + ( l 2 ) a i l − 2 s u m 2 + . . . + ( l l ) a i 0 s u m l ] x j = [ ( l 0 ) a i l d p i − 1 , j − 1 , 0 + ( l 1 ) a i l − 1 d p i − 1 , j − 1 , 1 + ( l 2 ) a i l − 2 d p i − 1 , j − 1 , 2 + . . . + ( l l ) a i 0 d p i − 1 , j − 1 , l ] x = [ 1 0 ! ( l − 0 ) ! a i l d p i − 1 , j − 1 , 0 + 1 1 ! ( l − 1 ) ! a i l − 1 d p i − 1 , j − 1 , 1 + 1 2 ! ( l − 2 ) ! a i l − 2 d p i − 1 , j − 1 , 2 + . . . + 1 l ! ( l − l ) ! a i 0 d p i − 1 , j − 1 , l ] x l ! \begin{aligned} dp_{i,j,l}&=[(p_0+a_i)^l+(p_1+a_i)^l+(p_2+a_i)^l+...]x^j\\ &=[\binom{l}{0} a_i^l sum_0+\binom{l}{1} a_i^{l-1} sum_1+\binom{l}{2} a_i^{l-2} sum_2+...+\binom{l}{l} a_i^0 sum_l]x^j\\ &=[\binom{l}{0} a_i^l dp_{i-1,j-1,0}+\binom{l}{1} a_i^{l-1} dp_{i-1,j-1,1}+\binom{l}{2} a_i^{l-2} dp_{i-1,j-1,2}+...+\binom{l}{l} a_i^0 dp_{i-1,j-1,l}]x\\ &=[\frac{1}{0!(l-0)!} a_i^l dp_{i-1,j-1,0}+\frac{1}{1!(l-1)!} a_i^{l-1} dp_{i-1,j-1,1}+\frac{1}{2!(l-2)!} a_i^{l-2} dp_{i-1,j-1,2}+...+\frac{1}{l!(l-l)!} a_i^0 dp_{i-1,j-1,l}]x l! \end{aligned} dpi,j,l=[(p0+ai)l+(p1+ai)l+(p2+ai)l+...]xj=[(0l)ailsum0+(1l)ail−1sum1+(2l)ail−2sum2+...+(ll)ai0suml]xj=[(0l)aildpi−1,j−1,0+(1l)ail−1dpi−1,j−1,1+(2l)ail−2dpi−1,j−1,2+...+(ll)ai0dpi−1,j−1,l]x=[0!(l−0)!1aildpi−1,j−1,0+1!(l−1)!1ail−1dpi−1,j−1,1+2!(l−2)!1ail−2dpi−1,j−1,2+...+l!(l−l)!1ai0dpi−1,j−1,l]xl!

我们发现化简到最后，参数 j j j并没有参与计算，因此可将式子转化为
d p i , l = [ 1 0 ! l ! a i l d p i − 1 , 0 + 1 1 ! ( l − 1 ) ! a i l − 1 d p i − 1 , 1 + 1 2 ! ( l − 2 ) ! a i l − 2 d p i − 1 , 2 + . . . + 1 l ! 0 ! a i 0 d p i − 1 , l ] x l ! dp_{i,l}=[\frac{1}{0!l!} a_i^l dp_{i-1,0}+\frac{1}{1!(l-1)!} a_i^{l-1} dp_{i-1,1}+\frac{1}{2!(l-2)!} a_i^{l-2} dp_{i-1,2}+...+\frac{1}{l!0!} a_i^0 dp_{i-1,l}]x l! dpi,l=[0!l!1aildpi−1,0+1!(l−1)!1ail−1dpi−1,1+2!(l−2)!1ail−2dpi−1,2+...+l!0!1ai0dpi−1,l]xl!

我们可以发现这就是个多项式乘法。

令 F ( u ) = 1 0 ! d p i − 1 , 0 u 0 + 1 1 ! d p i − 1 , 1 u 1 + 1 2 ! d p i − 1 , 2 u 2 + . . . + 1 l ! d p i − 1 , l u l F(u)=\frac{1}{0!} dp_{i-1,0}u^0+\frac{1}{1!} dp_{i-1,1}u^1+\frac{1}{2!} dp_{i-1,2}u^2+...+\frac{1}{l!} dp_{i-1,l}u^l F(u)=0!1dpi−1,0u0+1!1dpi−1,1u1+2!1dpi−1,2u2+...+l!1dpi−1,lul，

G ( u ) = 1 0 ! a i 0 u 0 + 1 1 ! a i 1 u 1 + 1 2 ! a i 2 u 2 + . . . + 1 l ! a i l u l G(u)=\frac{1}{0!} a_i^0 u^0+\frac{1}{1!} a_i^1 u^1+\frac{1}{2!} a_i^2 u^2+...+\frac{1}{l!} a_i^l u^l G(u)=0!1ai0u0+1!1ai1u1+2!1ai2u2+...+l!1ailul，

H ( u ) = F ( u ) G ( u ) x l ! H(u)=F(u)G(u)xl! H(u)=F(u)G(u)xl!。

那么 d p i , l = d p i − 1 , l + a i l x + [ u l ] H dp_{i,l}=dp_{i-1,l}+a_i^lx+[u^l]H dpi,l=dpi−1,l+ailx+[ul]H（ [ u l ] H [u^l]H [ul]H表示 H ( u ) H(u) H(u)的第 l l l项系数），最终答案为 d p n , k dp_{n,k} dpn,k。

#include<bits/stdc++.h>
#define ll long long
#define L(i,j,k) for(ll i=(j);i<=(k);i++)
#define R(i,j,k) for(ll i=(j);i>=(k);i--)
#define inf 0x3f3f3f3f3f3f3f3f
#define vec vector
#define pll pair<ll,ll>
#define fi first
#define se second
#define pb push_back
#define mkp make_pair
#define MS(i,j,k) memset(i,j,sizeof(ll)*(k+2))
const ll N=1e5+10,M=10;
const ll mod=998244353,mmod=mod-1;
const double pi=acos(-1),eps=1e-8;
using namespace std;
ll fmul(ll a,ll b){a%=mod;b%=mod;ll res=0;while(b){if(b&1){res+=a;res%=mod;}a<<=1;if(a>=mod)a%=mod;b>>=1;}return res;}
ll gcd(ll x,ll y){if(y==0) return x;return gcd(y,x%y);}
ll fksm(ll a,ll b){ll r=1;if(b<0)b+=mod-1;for(a%=mod;b;b>>=1){if(b&1)r=r*a%mod;a=a*a%mod;}return r;}//a 分母; b MOD-2
ll lowbit(ll x){return x&(-x);}ll m,n,t,x,y,z,l,r,u,v,k,p,pp,nx,ny,nz,ansx,ansy,num,sum,mn,mx,ans;
ll lim,pos,tot,cnt,root,key,block;
ll a[N],dp[N];ll fac[N],fra[N],two=fksm(2,mod-2);
void init(ll n){//n阶阶乘初始化 fac[0]=1;L(i,1,n)fac[i]=fac[i-1]*i%mod;fra[n]=fksm(fac[n],mod-2);R(i,n-1,0)fra[i]=fra[i+1]*(i+1)%mod;
}
ll C(ll n,ll k){if(!n&&!k)return 1;if(n<k||k<0)return 0;return fac[n]*fra[k]%mod*fra[n-k]%mod;}//组合数ll R[N],a0[N],a1[N];
void init_ntt(ll limit){lim=1;ll k=0;while(lim<limit)lim*=2,k++;L(i,0,lim-1)R[i]=(R[i>>1]>>1)|((i&1)<<(k-1));}
void ntt(ll *f,ll type){ll G=3,Gi=fksm(G,mod-2);//mod=998244353,G=3;mod=1e9+7,G=5;L(i,0,lim-1)if(i<R[i])swap(f[i],f[R[i]]);for(ll mid=1;mid<lim;mid<<=1){ll Wn=fksm(type==1?G:Gi,(mod-1)/(mid<<1));for(ll j=0;j<lim;j+=(mid<<1)){ll w=1;for(ll k=0;k<mid;k++,w=(w*Wn)%mod){ll x=f[j+k],y=w*f[j+k+mid]%mod;f[j+k]=(x+y)%mod,f[j+k+mid]=(x-y+mod)%mod;}}}
}void NTT(ll *f,ll *g,ll siz0,ll siz1){init_ntt(siz0+siz1-1);L(i,siz0,lim)f[i]=0;L(i,siz1,lim)g[i]=0;ntt(f,1);ntt(g,1);L(i,0,lim-1)f[i]=(f[i]*g[i])%mod;ntt(f,-1);ll inv=fksm(lim,mod-2);L(i,0,lim-1)f[i]=f[i]*inv%mod;
}void solve(){scanf("%lld%lld%lld",&n,&k,&m);MS(dp,0,k);L(i,1,n)scanf("%lld",&a[i]);L(i,1,n){L(j,0,k){a0[j]=dp[j]*fra[j]%mod;a1[j]=fksm(a[i],j)*fra[j]%mod;}NTT(a0,a1,k+1,k+1);L(j,0,k)dp[j]=(dp[j]+fksm(a[i],j)*m%mod+a0[j]*m%mod*fac[j]%mod)%mod;}printf("%lld\n",dp[k]);
}int main(){// ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);// cout<<fixed<<setprecision(12);//精度init(10010);ll T=1;scanf("%lld",&T);while(T--)solve();
}