模板 - 数学 - 多项式 - NTT

Posted 2020-11-23 kisekipurin2019

Huffman分治的NTT，常数一般。使用的时候把多项式的系数们放进vector里面，然后调用solve就可以得到它们的乘积。注意这里默认最大长度是1e6，可能需要改变。

#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
", vec[pq.top().id][pq.top().siz >> 1]);

    return 0;
}