Templates - Mathematics - polynomial - NTT

Huffman partition of NTT, generally constant. When using the coefficients of the polynomial into the vector were inside, and then you can get calling solve their product. Note that the default maximum length is 1e6, may need to change.

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;

int a[200005], b[200005], btop;

const int MAXN = 1e6, MAXLOGN = 20, mod = 998244353;

int add_mod(int x, int y) {
    x += y;
    if(x >= mod)
        x -= mod;
    return x;
}

int sub_mod(int x, int y) {
    x -= y;
    if(x < 0)
        x += mod;
    return x;
}

ll mul_mod(ll x, int y) {
    x *= y;
    if(x >= mod)
        x %= mod;
    return x;
}

int pow_mod(ll x, int n) {
    ll res = 1;
    while(n) {
        if(n & 1)
            res = mul_mod(res, x);
        x = mul_mod(x, x);
        n >>= 1;
    }
    return res;
}

int gl[MAXLOGN + 1];

void init() {
    for(int len = 2; len <= MAXN; len <<= 1)
        gl[__builtin_ctz(len)] = pow_mod(3, (mod - 1) / len);
}

void NTT(int a[], int n, int op) {
    for(int i = 1, j = n >> 1; i < n - 1; ++i) {
        if(i < j)
            swap(a[i], a[j]);
        int k = n >> 1;
        while(k <= j) {
            j -= k;
            k >>= 1;
        }
        j += k;
    }
    for(int len = 2; len <= n; len <<= 1) {
        int g = gl[__builtin_ctz(len)];
        for(int i = 0; i < n; i += len) {
            int w = 1;
            for(int j = i; j < i + (len >> 1); ++j) {
                int u = a[j], t = mul_mod(a[j + (len >> 1)], w);
                a[j] = add_mod(u, t), a[j + (len >> 1)] = sub_mod(u, t);
                w = mul_mod(w, g);
            }
        }
    }
    if(op == -1) {
        reverse(a + 1, a + n);
        int inv = pow_mod(n, mod - 2);
        for(int i = 0; i < n; ++i)
            a[i] = mul_mod(a[i], inv);
    }
}

int A[MAXN + 5], B[MAXN + 5];

int pow2(int x) {
    int res = 1;
    while(res < x)
        res <<= 1;
    return res;
}

void convolution(int A[], int B[], int Asize, int Bsize) {
    int n = pow2(Asize + Bsize - 1);
    memset(A + Asize, 0, sizeof(A[0]) * (n - Asize));
    memset(B + Bsize, 0, sizeof(B[0]) * (n - Bsize));
    NTT(A, n, 1);
    NTT(B, n, 1);
    for(int i = 0; i < n; ++i)
        A[i] = mul_mod(A[i], B[i]);
    NTT(A, n, -1);
    return;
}

vector<int> vec[200005], evec;
struct PriorityQueueNode {
    int siz, id;
    bool operator<(const PriorityQueueNode &pqn) const {
        return siz > pqn.siz;
    }
} pqn;

priority_queue<PriorityQueueNode> pq;

void solve() {
    //哈夫曼分治
    init();
    while(pq.size() > 1) {
        int Aid = pq.top().id, Asize = vec[Aid].size();
        for(int i = 0; i < Asize; ++i)
            A[i] = vec[Aid][i];
        pq.pop();

        int Bid = pq.top().id, Bsize = vec[Bid].size();
        for(int i = 0; i < Bsize; ++i)
            B[i] = vec[Bid][i];
        pq.pop();

        convolution(A, B, Asize, Bsize);
        Asize = Asize + Bsize - 1;

        vec[Aid].resize(Asize);
        for(int i = 0; i < Asize; ++i)
            vec[Aid][i] = A[i];
        pq.push({Asize, Aid});
        vec[Bid] = evec;
    }
}

int main() {
#ifdef KisekiPurin
    freopen("KisekiPurin.in", "r", stdin);
#endif // KisekiPurin
    int n;
    scanf("%d", &n);
    for(int i = 1; i <= n; ++i)
        scanf("%d", &a[i]);
    sort(a + 1, a + 1 + n);

    btop = 0;
    for(int i = 1; i <= n; ++i) {
        if(a[i] != a[i - 1])
            b[++btop] = 1;
        else
            ++b[btop];
    }

    sort(b + 1, b + 1 + btop);
    for(int i = 1; i <= btop; ++i) {
        while(vec[0].size() < b[i] + 1)
            vec[0].push_back(1);
        vec[i] = vec[0];
        pq.push({vec[i].size(), i});
    }

    solve();

    printf("%d\n", vec[pq.top().id][pq.top().siz >> 1]);

    return 0;
}