挑选队友 (生成函数 + FFT + 分治)

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

题目描述

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

输入描述:

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

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

n ≤ 100000, m ≤ k ≤ n

输出描述:

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

示例1

输入

5 3 4
1 2 2

输出

4

 

析：由于每个群都要选人，而且每个人还不同，从一个 n 个人的群里选 m 个人，方法数是 C(n, m)。但是要考虑多个群就是一个生成函数的问题了，答案就是[(x+1)^s1-1] * [(x+1)^s2-1] * ... * [(x+1)^sm-1]，该多项式的第 k +1 项的系数，也就是 x ^k 的系数。最多可能有 n 个多项式相乘，可能有 n 项，那么如果使用普通的 FFT 算法在时间上是过不去了，必须要使用分治来做进一步优化。

代码如下：

#pragma comment(linker, "/STACK:1024000000,1024000000")
#include <cstdio>
#include <string>
#include <cstdlib>
#include <cmath>
#include <iostream>
#include <cstring>
#include <set>
#include <queue>
#include <algorithm>
#include <vector>
#include <map>
#include <cctype>
#include <cmath>
#include <stack>
#include <sstream>
#include <list>
#include <assert.h>
#include <bitset>
#include <numeric>
#define debug() puts("++++")
#define gcd(a, b) __gcd(a, b)
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define fi first
#define se second
#define pb push_back
#define sqr(x) ((x)*(x))
#define ms(a,b) memset(a, b, sizeof a)
#define sz size()
#define be begin()
#define ed end()
#define pu push_up
#define pd push_down
#define cl clear()
#define lowbit(x) -x&x
//#define all 1,n,1
#define FOR(i,n,x)  for(int i = (x); i < (n); ++i)
#define freopenr freopen("in.in", "r", stdin)
#define freopenw freopen("out.out", "w", stdout)
using namespace std;
 
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> P;
const int INF = 0x3f3f3f3f;
const LL LNF = 1e17;
const double inf = 1e20;
const double PI = acos(-1.0);
const double eps = 1e-8;
const int maxn = 4e5 + 100;
const int maxm = 1e6 + 10;
const LL mod = 998244353LL;
const int dr[] = {-1, 1, 0, 0, 1, 1, -1, -1};
const int dc[] = {0, 0, 1, -1, 1, -1, 1, -1};
const char *de[] = {"0000", "0001", "0010", "0011", "0100", "0101", "0110", "0111", "1000", "1001", "1010", "1011", "1100", "1101", "1110", "1111"};
int n, m;
const int mon[] = {0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31};
const int monn[] = {0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31};
inline bool is_in(int r, int c) {
  return r >= 0 && r < n && c >= 0 && c < m;
}
inline int readInt(){ int x;  scanf("%d", &x);  return x; }
const int g = 3;
LL qp[40], fact[maxn>>2], inv[maxn>>2];

LL fast_pow(LL a, LL n){
  LL res = 1;
  while(n){
    if(n&1)  res = res * a % mod;
    n >>= 1;
    a = a * a % mod;
  }
  return res;
}

void init(){
  for(int i = 0; i < 30; ++i)
    qp[i] = fast_pow(g, (mod-1)/(1<<i));
  fact[0] = fact[1] = inv[0] = inv[1] = 1;
  for(int i = 2; i <= n; ++i)
    fact[i] = fact[i-1] * i % mod;
  inv[n] = fast_pow(fact[n], mod - 2);
  for(int i = n-1; i > 1; --i)  inv[i] = inv[i+1] * (i+1) % mod;
}

void rader(vector<LL> &F, int len){
  int j = len >> 1;
  for(int i = 1; i < len-1; ++i){
    if(i < j)  swap(F[i], F[j]);
    int k = len >> 1;
    while(j >= k){
      j ^= k;  k >>= 1;
    }
    if(j < k)  j |= k;
  }
}

void NTT(vector<LL> &F, int len, int t){
  int id = 0;
  rader(F, len);
  for(int h = 2; h <= len; h <<= 1){
    ++id;
    for(int j = 0; j < len; j += h){
      LL E = 1;
      for(int k = j; k < j+h/2; ++k){
        LL u = F[k];
        LL v = E * F[k+h/2] % mod;
        F[k] = (u + v) % mod;
        F[k+h/2] = (u - v + mod) % mod;
        E = E * qp[id] % mod;
      }
    }
  }
  if(t == -1){
    for(int i = 1; i < (len>>1); ++i)  swap(F[i], F[len-i]);
    LL inv = fast_pow(len, mod - 2);
    for(int i = 0; i < len; ++i)  F[i] = F[i] * inv % mod;
  }
}

inline LL C(int n, int m){
  return fact[n] * inv[m] % mod * inv[n-m] % mod;
}

vector<LL> solve(int len1, int len2, vector<LL> &A, vector<LL> &B){
  int len = 1;
  while(len < (len1<<1) || len < (len2<<1))  len <<= 1;
  A.resize(len);  B.resize(len);
  NTT(A, len, 1);
  NTT(B, len, 1);
  for(int i = 0; i < len; ++i)  A[i] = A[i] * B[i] % mod;
  NTT(A, len, -1);
  while(len > 1 && A[len-1] == 0)  --len;
  A.resize(len);
  return A;
}

vector<LL> sum[maxn];
vector<vector<LL> > all;

void dfs(int l, int r, int rt){
  if(l == r){ sum[rt] = all[l];  return ; }
  int m = l + r >> 1;
  dfs(lson);  dfs(rson);
  sum[rt] = solve(sum[rt<<1].sz, sum[rt<<1|1].sz, sum[rt<<1], sum[rt<<1|1]);
}


int main(){
  int k;
  scanf("%d %d %d", &n, &m, &k);
  init();  all.pb(vector<LL>());
  for(int i = 0; i < m; ++i){
    int x;  scanf("%d", &x);
    vector<LL> v;  v.pb(0);
    for(int j = 1; j <= x; ++j)  v.pb(C(x, j));
    all.pb(v);
  }
  dfs(1, m, 1);
  printf("%lld\n", sum[1][k]);
  return 0;
}