【清华集训2017】生成树计数 pufer序+分治NTT优化DP

Description

在一个$s$个点的图中，存在$s-n$条边，使图中形成了$n$个连通块，第$i$个连通块中有$a_i$个点。
现在我们需要再连接$n-1$条边，使该图变成一棵树。对一种连边方案，设原图中第$i$个连通块连出了$d_i$条边，那么这棵树$T$的价值为：
$$val(T)=(\prod _{i=1}^n{d}_i^m)(\sum _{i=1}^n{d}_i^m)$$你的任务是求出所有可能的生成树的价值之和，对$998244353$取模。

---

Sample Input

3 1
2 3 4

---

Sample Output

1728

---

首先一个$a_i$的贡献就是$a_i^{d_i}$。
先考虑这个式子的形式一个$sum$乘上一个乘积太难搞了。
我们把它变成有一个$i$可以多乘上一个$d_i^m$，那么你就可以设：
$f[i][j]$为$i$个点已经填了$pufer$序中的$j$个点的方案，
$g[i][j]$为$i$个点已经填了$pufer$序中的$j$个点，有了一个特殊点方案，
每次枚举当前这个点加进来多少点转移即可。

这个$d_i$枚举的范围太大了，不是很好搞，我们考虑一些组合意义来优化。
$d_i^m$相当于有$d_i$个格子，你可以染$m$次颜色，每次可以选任意一个格子涂颜色。但是你一个点的度数实际上是$d_i+1$，所以你每一个点都要在后面再额外预留一个格子。
你发现这样有颜色的格子是不超过$m$个的，状态改变为：$f[i][j]$表示$i$个颜色确定前半部分的$j$个格子，你枚举前半部分填了几个有颜色的格子，可得转移：
$$f[i][j]=\sum _{k=0}^{m}f[i-1][j-k]\ast C_{n-2-j+k}^k\ast a_i^{k+1}(S_m^k\ast k!+S_m^{k+1}\ast (k+1)!)$$
$g[i][j]$表示$i$个格子确定了$j$是有颜色的，已经确定了一个特殊点，可得转移：
$$g[i][j]=\sum _{k=0}^{m}g[i-1][j-k]\ast C_{n-2-j+k}^k\ast a_i^{k+1}(S_m^k\ast k!+S_m^{k+1}\ast (k+1)!)$$
$$g[i][j]+=\sum _{k=0}^{2m}f[i-1][j-k]\ast C_{n-2-j+k}^k\ast a_i^{k+1}(S_{2m}^k\ast k!+S_{2m}^{k+1}\ast (k+1)!)$$
那个斯特林数就是分别考虑后半部分填或不填。
最后再考虑把那些每个$a_i$占了但不涂颜色的位置都给乘回来，可得答案：
$$ans=\sum _{i=0}^{n-2}g[n][i]\ast (\sum _{j=1}^n{a}_j{)}^{n-2-i}$$
这个$dp$式直接用分治$NTT$优化即可做到$O(nmlo{g}^{2}n)$

---

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;
int _max(int x, int y) {return x > y ? x : y;}
int _min(int x, int y) {return x < y ? x : y;}
const int mod = 998244353;
const int N = 30001, M = 31;
int read() {
	int s = 0, f = 1; char ch = getchar();
	while(ch < '0' || ch > '9') {if(ch == '-') f = -1; ch = getchar();}
	while(ch >= '0' && ch <= '9') s = s * 10 + ch - '0', ch = getchar();
	return s * f;
}
void put(int x) {
	if(x >= 10) put(x / 10);
	putchar(x % 10 + '0');
}

int jc[N], inv[N];
int n, m, R[N << 2], S[M * 2][M * 2], d1[N][M * 2], d2[N][M * 2];
int a[N], f[20][N << 2], g[20][N << 2];

int pow_mod(int a, int k) {
	int ans = 1;
	while(k) {
		if(k & 1) ans = (LL)ans * a % mod;
		a = (LL)a * a % mod, k /= 2;
	} return ans;
}

int add(int x, int y) {
	x += y;
	return x >= mod ? x - mod : x;
}

int dec(int x, int y) {
	x -= y;
	return x < 0 ? x + mod : x;
}

void NTT(int y[], int len, int on) {
	for(int i = 0; i < len; i++) if(i < R[i]) swap(y[i], y[R[i]]);
	for(int i = 1; i < len; i <<= 1) {
		int wn = pow_mod(3, (mod - 1) / (i << 1));
		if(on == -1) wn = pow_mod(wn, mod - 2);
		for(int j = 0; j < len; j += i << 1) {
			int w = 1;
			for(int k = 0; k < i; k++, w = (LL)w * wn % mod) {
				int u = y[j + k], v = (LL)y[j + k + i] * w % mod;
				y[j + k] = add(u, v), y[j + k + i] = dec(u, v);
			}
		}
	} if(on == -1) {
		int hh = pow_mod(len, mod - 2);
		for(int i = 0; i < len; i++) y[i] = (LL)y[i] * hh % mod;
	}
}

void pre() {
	jc[0] = inv[0] = 1; for(int i = 1; i <= n; i++) jc[i] = (LL)jc[i - 1] * i % mod;
	inv[n] = pow_mod(jc[n], mod - 2); for(int i = n - 1; i >= 1; i--) inv[i] = (LL)inv[i + 1] * (i + 1) % mod;
	S[0][0] = 1;
	for(int i = 1; i <= 2 * m; i++) {
		for(int j = 1; j <= _min(n - 1, 2 * m); j++) S[i][j] = add(S[i - 1][j - 1], (LL)S[i - 1][j] * j % mod);
	} for(int i = 1; i <= n; i++) {
		d1[i][0] = d2[i][0] = 1;
		for(int j = 1; j <= _min(n - 2, m); j++) d1[i][j] = (LL)d1[i][j - 1] * a[i] % mod;
		for(int j = 1; j <= _min(n - 2, 2 * m); j++) d2[i][j] = (LL)d2[i][j - 1] * a[i] % mod;
	} for(int i = 1; i <= n; i++) {
		for(int j = 0; j <= _min(n - 2, m); j++) d1[i][j] = (LL)d1[i][j] * inv[j] % mod * add((LL)S[m][j] * jc[j] % mod, (LL)S[m][j + 1] * jc[j + 1] % mod) % mod;
		for(int j = 0; j <= _min(n - 2, 2 * m); j++) d2[i][j] = (LL)d2[i][j] * inv[j] % mod * add((LL)S[2 * m][j] * jc[j] % mod, (LL)S[2 * m][j + 1] * jc[j + 1] % mod) % mod;
	}
}

void solve(int l, int r, int d) {
	if(l == r) {
		for(int i = 0; i <= _min(n, 2 * m); i++) g[d][i] = d2[l][i], f[d][i] = d1[l][i];
		return ;
	} int mid = (l + r) / 2;
	int L1 = _min((mid - l + 1) * m + m, n - 2), L2 = _min((r - mid) * m + m, n - 2), L;
	for(L = 1; L <= (L1 + L2); L <<= 1);
	solve(l, mid, d + 1);
	for(int i = 0; i <= L1; i++) f[d][i] = f[d + 1][i];
	for(int i = 0; i <= L1; i++) g[d][i] = g[d + 1][i];
	solve(mid + 1, r, d + 1);
	for(int i = 0; i < L; i++) R[i] = (R[i >> 1] >> 1) | ((i & 1) * (L >> 1));
	for(int i = L1 + 1; i < L; i++) f[d][i] = g[d][i] = 0;
	for(int i = L2 + 1; i < L; i++) f[d + 1][i] = g[d + 1][i] = 0;
	NTT(f[d], L, 1), NTT(g[d], L, 1), NTT(f[d + 1], L, 1), NTT(g[d + 1], L, 1);
	for(int i = 0; i < L; i++) g[d][i] = add((LL)f[d + 1][i] * g[d][i] % mod, (LL)f[d][i] * g[d + 1][i] % mod), f[d][i] = (LL)f[d][i] * f[d + 1][i] % mod;
	NTT(f[d], L, -1), NTT(g[d], L, -1);
	int uu = (r - l + 1) * m + m;
	for(int i = _min(uu, n - 2) + 1; i < L; i++) g[d][i] = 0;
	for(int i = _min(uu, n - 2) + 1; i < L; i++) f[d][i] = 0;
}

int main() {
	n = read(), m = read(); int sum = 0;
	for(int i = 1; i <= n; i++) a[i] = read(), sum = add(sum, a[i]);
	pre(), solve(1, n, 0); int ans = 0, o = 1;
	for(int i = n - 2; i >= 0; i--) {
		ans = add(ans, (LL)g[0][i] * inv[n - 2 - i] % mod * o % mod);
		o = (LL)o * sum % mod;
	} for(int i = 1; i <= n; i++) ans = (LL)ans * a[i] % mod;
	put((LL)ans * jc[n - 2] % mod);
	return 0;          
}