多项式相乘NTT模板题——Wannafly挑战赛20D.挑选队友

链接：https://ac.nowcoder.com/acm/contest/133/D
来源：牛客网
 

时间限制：C/C++ 1秒，其他语言2秒
空间限制：C/C++ 262144K，其他语言524288K
64bit IO Format: %lld

题目描述

Applese打开了m个QQ群，向群友们发出了组队的邀请。作为网红选手，Applese得到了n位选手的反馈，每位选手只会在一个群给Applese反馈
现在，Applese要挑选其中的k名选手组队比赛，为了维持和各个群的良好关系，每个群中都应有至少一名选手成为Applese的队友（数据保证每个群都有选手给Applese反馈）
Applese想知道，他有多少种挑选队友的方案

输入描述:

输入包括两行
第一行包括三个数n, m, k，表示共有n位选手，m个群，需要有k名选手被选择

第二行包括m个数，第i个数表示第i个群有si个选手

n ≤ 100000, m ≤ k ≤ n
∑i=1msi=n\sum_{i=1}^m s_i = n∑i=1m​si​=n
∀i,si>0\forall i, s_i > 0∀i,si​>0

输出描述:

输出包括一行
第一行输出方案数
由于输出可能比较大，你只需要输出在模998244353意义下的答案

示例1

输入

复制

5 3 4
1 2 2

输出

复制

4

利用普通型生成函数的定理，将同一组的人视为同一个变量ai，那么现在ai的限制条件就是至少有一个，所以对于每一个分组，假如他的大小为si，那么生成函数为

然后有m个这样的多项式相乘，那就变成NTT模板题了。

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
#include <queue>
using namespace std;
typedef long long ll;
const ll mod = 998244353;
const int maxn = 3e5 + 111;
inline ll qpow(ll a, ll b)
{
	ll sum = 1;
	while (b)
	{
		if (b & 1)
			sum = sum * a % mod;
		b >>= 1;
		a = a * a % mod;
	}
	return sum;
}
inline ll Inv(ll a, ll _mod)
{
	return qpow(a, _mod - 2);
}
struct NTT
{
	int rev[maxn], dig[105];
	int N, L;
	ll g;
	void init_rev(int n)
	{
		//初始化原根
		g = 3;
		for (N = 1, L = 0; N <= n; N <<= 1, L++);
		memset(dig, 0, sizeof(int) * (L + 1));
		for (int i = 0; i < N; i++)
		{
			rev[i] = 0;
			int len = 0;
			for (int t = i; t; t >>= 1)
				dig[len++] = t & 1;
			for (int j = 0; j < L; j++)
				rev[i] = (rev[i] << 1) | dig[j];
		}
	}

	void DFT(vector<ll>& a, int flag)
	{
		for (int i = 0; i < N; i++)
			if (i < rev[i])
				swap(a[i], a[rev[i]]);

		for (int l = 1; l < N; l <<= 1)
		{
			ll wn;
			if (flag == 1)
				wn = qpow(g, (mod - 1) / (2 * l));
			else
				wn = qpow(g, mod - 1 - (mod - 1) / (2 * l));
			for (int k = 0; k < N; k += l * 2)
			{
				ll w = 1;
				ll x, y;
				for (int j = k; j < k + l; j++)
				{
					x = a[j];
					y = a[j + l] * w % mod;
					a[j] = (x + y) % mod;
					a[j + l] = (x - y + mod) % mod;
					w = w * wn % mod;
				}
			}
		}
		if (flag == -1)
		{
			ll x = Inv(N, mod);
			for (int i = 0; i < N; i++)
				a[i] = a[i] * x % mod;
		}
	}

	void mul(vector<ll>& a, vector<ll>& b, int m)
	{
		init_rev(m);
		a.resize(N);
		b.resize(N);
		DFT(a, 1);
		DFT(b, 1);
		for (int i = 0; i < N; i++)
			a[i] = a[i] * b[i] % mod;
		DFT(a, -1);
		int len = N - 1;
		while (a[len] == 0) len--;
		a.resize(len + 1);
	}
}ntt;
vector<ll> v[maxn];
ll fac[maxn], inv[maxn];
void init(int n)
{
	fac[0] = fac[1] = inv[0] = inv[1] = 1;
	for (int i = 2; i <= n; i++)
		fac[i] = fac[i - 1] * i % mod;
	inv[n] = Inv(fac[n], mod);
	for (int i = n - 1; i > 1; i--)
		inv[i] = inv[i + 1] * (i + 1) % mod;
}
inline ll Comb(int n, int m)
{
	return fac[n] * inv[n - m] % mod * inv[m] % mod;
}
signed main() {
	int n, m, k;
	scanf("%d%d%d", &n, &m, &k);
	init(n + 1);
	queue<int>Q;
	for (int i = 1; i <= m; i++) {
		int x;
		scanf("%d", &x);
		for (int j = 1; j <= x; j++) {
			v[i].push_back(Comb(x, j));
		}
		Q.push(i);
	}
	while (Q.size() > 1) {
		int q1 = Q.front();
		Q.pop();
		int q2 = Q.front();
		Q.pop();
		ntt.mul(v[q1], v[q2], v[q1].size() + v[q2].size());
		Q.push(q1);
	}
	int top = Q.front();
	ll ans = v[top][k - m];
	printf("%lld\n", ans);
	return 0;
}