Hackerrank--Emma and sum of products (FFT)

题目链接

Emma is really fond of integers and loves playing with them. Her friends were jealous, and to test her, one of them gave her a problem. 
Emma is given a list A of N integers and is asked a set of Q queries. Each query is denoted by an integer K, for which you have to return the sum of product of all possible sublists having exactly K elements. 
Emma has got stuck in this problem and you being her best friend have decided to help her write a code to solve it. Since the answers can be very large, print the answers modulo100003.

Input Format
First line has an integer N, denoting the number of integers in list A. Next line contains N space separated integers. The third line contains integer Q, and next Q lines have a single integer K.

Output Format
For each of the queries, print the corresponding answer in a new line.

NOTE Sublist here refers to selecting K elements from a list of N elements. There will be (NK) ways to do that, it doesn't matter if two elements are same.

Constraints
1≤N≤3×104 
1≤Ai≤105 
1≤Q≤N 
1≤K≤N

Sample Input #00

3
1 2 3
2
1
2

Sample Output #00

6
11

Sample Input #01

3
1 2 2
1
2

Sample Output #01

8

Explanation

Sample #00: 
For K=1 possible sublists are {1},{2},{3} so answer is 1+2+3=6. 
For K=2 possible sublists are {1,2},{2,3},{3,1} so answer is (1×2)+(2×3)+(3×1)=2+6+3=11.

Sample #01: 
For K=2 possible sublists are {1,2},{2,2},{2,1} so answer is (1×2)+(2×2)+(2×1)=2+4+2=8.

题意：给出n个数，有q次询问，每次询问一个k，从n个数中选出k个数，对这k个数做乘积，求所有可能的选法的和。

我们把每个数看做一个多项式：x + A[i]，　那么就可以得到ｎ个这样的多项式。　将这ｎ个多项式相乘，

那么ｋ对应的查询的答案就是多项式中x^k项的系数，和母函数有点类似。

下面就是如何计算ｎ个多项式的乘积，想到用ＦＦＴ，但是不可以直接线性的计算这ｎ个多项式的乘积，时间复杂度

太高，所以想到分治的思想，总的复杂度就是Ｏ(n*(log(n)^2))

Accepted Code:

#define _CRT_SECURE_NO_WARNINGS
#include <string>
#include <vector>
#include <algorithm>
#include <numeric>
#include <set>
#include <map>
#include <queue>
#include <iostream>
#include <sstream>
#include <cstdio>
#include <cmath>
#include <ctime>
#include <cstring>
#include <cctype>
#include <cassert>
#include <limits>
#include <bitset>
#include <complex>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define reu(i,l,u) for(int (i)=(int)(l);(i)<(int)(u);++(i))
#define all(o) (o).begin(), (o).end()
#define pb(x) push_back(x)
#define mp(x,y) make_pair((x),(y))
#define mset(m,v) memset(m,v,sizeof(m))
#define INF 0x3f3f3f3f
#define INFL 0x3f3f3f3f3f3f3f3fLL
using namespace std;
typedef vector<int> vi; typedef pair<int,int> pii; typedef vector<pair<int,int> > vpii;
typedef long long ll; typedef vector<long long> vl; typedef pair<long long,long long> pll; typedef vector<pair<long long,long long> > vpll;
typedef vector<string> vs; typedef long double ld;
template<typename T, typename U> inline void amin(T &x, U y) { if(y < x) x = y; }
template<typename T, typename U> inline void amax(T &x, U y) { if(x < y) x = y; }

typedef long double Num;    //??????long double?????
const Num PI = 3.141592653589793238462643383279L;
typedef complex<Num> Complex;
//n?????
//a?????
void fft_main(int n, Num theta, Complex a[]) {
    for(int m = n; m >= 2; m >>= 1) {
        int mh = m >> 1;
        Complex thetaI = Complex(0, theta);
        rep(i, mh) {
            Complex w = exp((Num)i*thetaI);
            for(int j = i; j < n; j += m) {
                int k = j + mh;
                Complex x = a[j] - a[k];
                a[j] += a[k];
                a[k] = w * x;
            }
        }
        theta *= 2;
    }
    int i = 0;
    reu(j, 1, n-1) {
        for(int k = n >> 1; k > (i ^= k); k >>= 1) ;
        if(j < i) swap(a[i], a[j]);
    }
}

void fft(int n, Complex a[]) { fft_main(n, 2 * PI / n, a); }
void inverse_fft(int n, Complex a[]) { fft_main(n, -2 * PI / n, a); }

void convolution(vector<Complex> &v, vector<Complex> &w) {
    int n = 1, vwn = v.size() + w.size() - 1;
    while(n < vwn) n <<= 1;
    v.resize(n), w.resize(n);
    fft(n, &v[0]);
    fft(n, &w[0]);
    rep(i, n) v[i] *= w[i];
    inverse_fft(n, &v[0]);
    rep(i, n) v[i] /= n;
}

const int MOD = 100003;

vector<int> calc_dfs(const vector<int> &A, int l, int r) {
    if(r - l == 1) {
        vector<int> res(2);
        res[0] = 1;
        res[1] = A[l];
        return res;
    }
    int mid = (l + r) / 2;
    vector<int> L = calc_dfs(A, l, mid), R = calc_dfs(A, mid, r);
    vector<Complex> Lc(all(L)), Rc(all(R));
    convolution(Lc, Rc);
    int n = L.size() + R.size() - 1;
    vector<int> res(n);
    rep(i, n) res[i] = (long long)(Lc[i].real() + .5) % MOD;
//    cerr << "["<< l << "," << r <<"): ";
//    rep(i, n) cerr << res[i] << ", "; cerr << endl;
    return res;
}

int main() {
    int N;
    scanf("%d", &N);
    vector<int> A(N);
    rep(i, N) {
        scanf("%d", &A[i]);
        A[i] %= MOD;
    }
    vector<int> ans = calc_dfs(A, 0, N);
    int Q;
    scanf("%d", &Q);
    rep(ii, Q) {
        int K;
        scanf("%d", &K);
        printf("%d\n", ans[K]);
    }
    return 0;
}