Loj 2320.「清华集训 2017」生成树计数
Loj 2320.「清华集训 2017」生成树计数
题目描述
在一个 \(s\) 个点的图中，存在 \(s-n\) 条边，使图中形成了 \(n\) 个连通块，第 \(i\) 个连通块中有 \(a_i\) 个点。
现在我们需要再连接 \(n-1\) 条边，使该图变成一棵树。对一种连边方案，设原图中第 \(i\) 个连通块连出了 \(d_i\) 条边，那么这棵树 \(T\) 的价值为：
\[\mathrm{val}(T) = \left(\prod_{i=1}^{n} {d_i}^m\right)\left(\sum_{i=1}^{n} {d_i}^m\right)
\]
你的任务是求出所有可能的生成树的价值之和，对 \(998244353\) 取模。
输入格式
输入的第一行包含两个整数 \(n,m\)，意义见题目描述。
接下来一行有 \(n\) 个整数，第 \(i\) 个整数表示 \(a_i\) \((1\le a_i< 998244353)\)。
* 你可以由 \(a_i\) 计算出图的总点数 \(s\)，所以在输入中不再给出 \(s\) 的值。
输出格式
输出包含一行一个整数，表示答案。
数据范围与提示
本题共有 \(20\) 个测试点，每个测试点 \(5\) 分。
- \(20\%\) 的数据中，\(n\le500\)。
- 另外 \(20\%\) 的数据中，\(n \le 3000\)。
- 另外 \(10\%\) 的数据中，\(n \le 10010, m = 1\)。
- 另外 \(10\%\)的数据中，\(n \le 10015,m = 2\)。
-另外 \(20\%\) 的数据中，所有 \(a_i\) 相等。
\(\\\)
好神的题啊！
假设我们知道了每个点的度数，考虑计算此时的生成树的个数。这个用\(prufer\)序列非常好解决：
假设第\(i\)个点在\(prufer\)序列中出现次数为\(d_i\),（则其度数为\(d_i+1\)）
\[Ans=(n-2)!\prod_{i=1}^n\frac{{a_i}^{d_i+1}}{d_i!}
\]
先考虑对式子进行变形
\[\begin{align}
Ans&=
\sum_{\sum d_i==n-2}
(n-2)!
\sum_{i=1}^n\frac{{{a_i}^{d_i+1}d_i}^{2m}}{d_i!}
\prod_{j=1,j\neq i}^n\frac{{d_j}^m}{d_j!}\\
&=(n-2)!\prod_{i=1}^na_i
\sum_{\sum_{d_i==n-2}}\sum_{i=1}^n\frac{{{a_i}^{d_i}d_i}^{2m}}{d_i!}
\prod_{j=1,j\neq i}^n\frac{{d_j}^m}{d_j!}\\
\end{align}
\]
设
\[Ans'=\sum_{\sum_{d_i==n-2}}\sum_{i=1}^n\frac{{{a_i}^{d_i}d_i}^{2m}}{d_i!}\prod_{j=1,j\neq i}^n\frac{{d_j}^m}{d_j!}
\]
考虑用生成函数解决：
\[A(x)=\sum_{i=0}^n\frac{i^{2m}}{i!}x^i\\
B(x)=\sum_{i=0}^n\frac{i^m}{i!}x^i
\]
则\(Ans'\)的生成函数为
\[\sum_{i=1}^nA(a_i)\prod_{j=1,j\neq i}^nB(a_j)\\
=\sum_{i=1}^n\frac{A(a_i)}{B(a_i)}\prod_{j=1}^nB(a_j)
\]
对于\(\prod_{j=1}^nB(a_j)\)，我们的一般套路是将其写成
\[\exp(\ln(\prod_{j=1}^nB(a_j)))\\
=\exp(\sum_{j=1}^n\ln(B(a_j)))
\]
这样做的好处是我们只需要求出\(\ln(B(x))\)，然后对第\(i\)项系数乘上\(\displaystyle \sum_{j=1}^n{a_j}^i\)就可以得到\(\displaystyle \sum_{j=1}^n\ln(B(a_j))\)了。对于\(\displaystyle \sum_{i=1}^n\frac{A(a_i)}{B(a_i)}\)我们也 用相同的处理方式。
所以:
\[Ans'=\sum_{i=1}^n\frac{A}{B}(a_i)\exp(\sum_{j=1}^n\ln(B(a_j)))
\]
现在的问题是如何求出
\[\sum_{i=1}^n{a_i}^k
\]
考虑\(\ln(x)\)的取\(x_0=1\)时的泰勒展开形式
\[\ln(x)=\sum_{i=0}\frac{\ln^{[i](1)}}{i!}(x-1)^i\\
=\sum_{i=1}\frac{(-1)^{i-1}}{i}(x-1)^i
\]
所以:
\[\ln(1+a_jx)=\sum_{i=1}\frac{(-1)^{i-1}{a_j}^i}{i}x^i
\]
那么我们只需要求出
\[\sum_{i=1}^n\ln(a_i)
\]
就行了。
\[\sum_{i=1}^n\ln(a_i)=\ln(\prod_{i=1}^n(1+a_ix))
\]
\(\prod_{i=1}^n(1+a_ix)\)可以用分治\(NTT\)求出。
代码：
#include<bits/stdc++.h>
#define ll long long
#define N 200005
using namespace std;
inline int Get() {int x=0,f=1;char ch=getchar();while(ch<'0'||ch>'9') {if(ch=='-') f=-1;ch=getchar();}while('0'<=ch&&ch<='9') {x=(x<<1)+(x<<3)+ch-'0';ch=getchar();}return x*f;}
const ll mod=998244353;
ll ksm(ll t,ll x) {
ll ans=1;
for(;x;x>>=1,t=t*t%mod)
if(x&1) ans=ans*t%mod;
return ans;
}
int n,m;
int a[N];
void NTT(ll *a,int d,int flag) {
int n=1<<d;
static int rev[N<<2];
static ll G=3;
for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<d-1);
for(int i=0;i<n;i++) if(i<rev[i]) swap(a[i],a[rev[i]]);
for(int s=1;s<=d;s++) {
int len=1<<s,mid=len>>1;
ll w=flag==1?ksm(G,(mod-1)/len):ksm(G,mod-1-(mod-1)/len);
for(int i=0;i<n;i+=len) {
ll t=1;
for(int j=0;j<mid;j++,t=t*w%mod) {
ll u=a[i+j],v=a[i+j+mid]*t%mod;
a[i+j]=(u+v)%mod;
a[i+j+mid]=(u-v+mod)%mod;
}
}
}
if(flag==-1) {
ll inv=ksm(n,mod-2);
for(int i=0;i<n;i++) a[i]=a[i]*inv%mod;
}
}
ll A[N<<2],B[N<<2];
ll inv[N<<2];
ll f[N<<2],g[N<<2];
void Inv(ll *inv,ll *a,int d) {
static ll A[N<<3];
if(d==0) {
inv[0]=ksm(a[0],mod-2);
return ;
}
Inv(inv,a,d-1);
for(int i=0;i<1<<d;i++) A[i]=a[i];
for(int i=1<<d;i<1<<d+1;i++) inv[i]=A[i]=0;
NTT(A,d+1,1);
NTT(inv,d+1,1);
for(int i=0;i<1<<d+1;i++) {
inv[i]=(2*inv[i]-A[i]*inv[i]%mod*inv[i]%mod+mod)%mod;
}
NTT(inv,d+1,-1);
for(int i=1<<d;i<1<<d+1;i++) inv[i]=0;
}
void Der(ll *a,int d) {
int n=1<<d;
for(int i=0;i<n-1;i++) a[i]=(i+1)*a[i+1]%mod;
a[n-1]=0;
}
void Int(ll *a,int d) {
int n=1<<d;
for(int i=n-1;i>0;i--) a[i]=ksm(i,mod-2)*a[i-1]%mod;
a[0]=0;
}
ll ln[N<<2];
void Ln(ll *ln,ll *a,int d) {
static ll der[N<<2];
for(int i=0;i<1<<d+1;i++) der[i]=0;
for(int i=0;i<1<<d;i++) der[i]=a[i];
Inv(inv,a,d);
Der(der,d);
NTT(inv,d+1,1),NTT(der,d+1,1);
for(int i=0;i<1<<d+1;i++) ln[i]=inv[i]*der[i]%mod;
NTT(ln,d+1,-1);
for(int i=1<<d;i<1<<d+1;i++) ln[i]=0;
Int(ln,d);
for(int i=1<<d;i<1<<d+1;i++) ln[i]=0;
}
ll ex[N<<2];
void Exp(ll *exp,ll *a,int d) {
static ll A[N<<2],B[N<<2];
if(d==0) {
exp[0]=1;
return ;
}
Exp(exp,a,d-1);
for(int i=0;i<1<<d;i++) A[i]=a[i];
for(int i=1<<d;i<1<<d+1;i++) exp[i]=A[i]=0;
Ln(B,exp,d);
NTT(exp,d+1,1);
NTT(B,d+1,1);
NTT(A,d+1,1);
for(int i=0;i<1<<d+1;i++) {
exp[i]=exp[i]*(1-B[i]+A[i]+mod)%mod;
}
NTT(exp,d+1,-1);
for(int i=1<<d;i<1<<d+1;i++) exp[i]=0;
}
void solve(int l,int r,ll *a) {
static ll A[N<<2],B[N<<2];
if(l==r) return ;
int mid=l+r>>1;
solve(l,mid,a),solve(mid+1,r,a);
int d=ceil(log2(r-l+2));
for(int i=0;i<1<<d;i++) A[i]=B[i]=0;
for(int i=l;i<=mid;i++) A[i-l+1]=a[i];
for(int i=mid+1;i<=r;i++) B[i-mid]=a[i];
A[0]=B[0]=1;
NTT(A,d,1),NTT(B,d,1);
for(int i=0;i<1<<d;i++) A[i]=A[i]*B[i]%mod;
NTT(A,d,-1);
for(int i=l;i<=r;i++) a[i]=A[i-l+1];
}
ll summ[N];
ll cal(int k) {
ll ans=0;
for(int i=1;i<=n;i++) (ans+=ksm(a[i],k))%=mod;
return ans;
}
ll tem[N<<2];
ll fac[N],ifac[N];
int main() {
n=Get(),m=Get();
fac[0]=1;
for(int i=1;i<=n;i++) fac[i]=fac[i-1]*i%mod;
ifac[n]=ksm(fac[n],mod-2);
for(int i=n-1;i>=0;i--) ifac[i]=ifac[i+1]*(i+1)%mod;
for(int i=1;i<=n;i++) a[i]=Get();
for(int i=1;i<=n;i++) summ[i]=a[i];
int d=ceil(log2(2*n+1));
solve(1,n,summ);
summ[0]=1;
Ln(ln,summ,d);
memcpy(summ,ln,sizeof(summ));
summ[0]=n;
for(int i=1;i<=n;i++) {
if(!(i&1)) summ[i]=summ[i]*(mod-1)%mod;
summ[i]=summ[i]*i%mod;
}
for(int i=0;i<=n;i++) {
A[i]=ksm(i+1,2*m)*ifac[i]%mod;
B[i]=ksm(i+1,m)*ifac[i]%mod;
}
Ln(ln,B,d);
for(int i=0;i<1<<d;i++) ln[i]=ln[i]*summ[i]%mod;
for(int i=n;i<1<<d;i++) ln[i]=0;
Exp(g,ln,d);
for(int i=n;i<=1<<d;i++) g[i]=0;
Inv(inv,B,d);
for(int i=n;i<1<<d;i++) inv[i]=0;
NTT(inv,d,1),NTT(A,d,1);
for(int i=0;i<1<<d;i++) f[i]=inv[i]*A[i]%mod;
NTT(f,d,-1);
for(int i=0;i<1<<d;i++) f[i]=f[i]*summ[i]%mod;
for(int i=n;i<1<<d;i++) f[i]=0;
NTT(f,d,1),NTT(g,d,1);
for(int i=0;i<1<<d;i++) f[i]=f[i]*g[i]%mod;
NTT(f,d,-1);
ll ans=fac[n-2];
for(int i=1;i<=n;i++) ans=ans*a[i]%mod;
ans=ans*f[n-2]%mod;
cout<<ans;
return 0;
}
