挑选队友 (生成函数 + 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个选手

n ≤ 100000, m ≤ k ≤ n

输出描述:

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

输入例子:

5 3 4

1 2 2

输出例子:

4

-->

示例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;

}