NTT (Convolution)

O(n log n)

Polynomial multiplication under the NTT-friendly prime 998244353 (primitive root 3): FFT over the integers mod p, so there is no floating-point error. Multiplying two polynomials, counting pair sums, and string wildcard matching all reduce to this convolution.

// use it on

AtCoder Library Practice F

Convolution of two sequences mod 998244353, which is the multiply function verbatim.

// the code

const long long MOD = 998244353, G = 3;  // MOD prime, G primitive root

long long power(long long b, long long e) {
    long long r = 1;
    for (b %= MOD; e > 0; e >>= 1, b = b * b % MOD)
        if (e & 1) r = r * b % MOD;
    return r;
}

void ntt(vector<long long>& a, bool invert) {
    int n = a.size();
    for (int i = 1, j = 0; i < n; i++) {  // bit-reversal permutation
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) swap(a[i], a[j]);
    }
    for (int len = 2; len <= n; len <<= 1) {
        long long w = power(G, (MOD - 1) / len);
        if (invert) w = power(w, MOD - 2);
        for (int i = 0; i < n; i += len) {
            long long wn = 1;
            for (int j = 0; j < len / 2; j++) {
                long long u = a[i + j];
                long long v = a[i + j + len / 2] * wn % MOD;
                a[i + j] = (u + v) % MOD;
                a[i + j + len / 2] = (u - v + MOD) % MOD;
                wn = wn * w % MOD;
            }
        }
    }
    if (invert) {
        long long ninv = power(n, MOD - 2);
        for (auto& x : a) x = x * ninv % MOD;
    }
}

vector<long long> multiply(vector<long long> a, vector<long long> b) {
    int res_size = a.size() + b.size() - 1, n = 1;
    while (n < res_size) n <<= 1;
    a.resize(n);
    b.resize(n);
    ntt(a, false);
    ntt(b, false);
    for (int i = 0; i < n; i++) a[i] = a[i] * b[i] % MOD;
    ntt(a, true);
    a.resize(res_size);
    return a;
}