Math/tests/ntt_any_mod.test.cpp

• View this file on GitHub
• Last update: 2023-10-15 09:43:20+08:00
• Problem: https://judge.yosupo.jp/problem/convolution_mod_1000000007

Depends on

• Math/Polynomial/NTT.h
• Math/modint.h

Code

#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod_1000000007"

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

#include "../modint.h"
#include "../Polynomial/NTT.h"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)
using mint = ModInt<1'000'000'007>;

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n, m; cin >> n >> m;
    vector<mint> a(n); REP(i,n) cin >> a[i];
    vector<mint> b(m); REP(i,m) cin >> b[i];

    auto c = multiply(a, b);
    for (const auto& val : c) cout << val << ' ';
    cout << endl;
    return 0;
}

#line 1 "Math/tests/ntt_any_mod.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/convolution_mod_1000000007"

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

#line 1 "Math/modint.h"
// ModInt {{{
template<int MD> struct ModInt {
    using ll = long long;
    int x;

    constexpr ModInt() : x(0) {}
    constexpr ModInt(ll v) { _set(v % MD + MD); }
    constexpr static int mod() { return MD; }
    constexpr explicit operator bool() const { return x != 0; }

    constexpr ModInt operator + (const ModInt& a) const {
        return ModInt()._set((ll) x + a.x);
    }
    constexpr ModInt operator - (const ModInt& a) const {
        return ModInt()._set((ll) x - a.x + MD);
    }
    constexpr ModInt operator * (const ModInt& a) const {
        return ModInt()._set((ll) x * a.x % MD);
    }
    constexpr ModInt operator / (const ModInt& a) const {
        return ModInt()._set((ll) x * a.inv().x % MD);
    }
    constexpr ModInt operator - () const {
        return ModInt()._set(MD - x);
    }

    constexpr ModInt& operator += (const ModInt& a) { return *this = *this + a; }
    constexpr ModInt& operator -= (const ModInt& a) { return *this = *this - a; }
    constexpr ModInt& operator *= (const ModInt& a) { return *this = *this * a; }
    constexpr ModInt& operator /= (const ModInt& a) { return *this = *this / a; }

    friend constexpr ModInt operator + (ll a, const ModInt& b) {
        return ModInt()._set(a % MD + b.x);
    }
    friend constexpr ModInt operator - (ll a, const ModInt& b) {
        return ModInt()._set(a % MD - b.x + MD);
    }
    friend constexpr ModInt operator * (ll a, const ModInt& b) {
        return ModInt()._set(a % MD * b.x % MD);
    }
    friend constexpr ModInt operator / (ll a, const ModInt& b) {
        return ModInt()._set(a % MD * b.inv().x % MD);
    }

    constexpr bool operator == (const ModInt& a) const { return x == a.x; }
    constexpr bool operator != (const ModInt& a) const { return x != a.x; }

    friend std::istream& operator >> (std::istream& is, ModInt& other) {
        ll val; is >> val;
        other = ModInt(val);
        return is;
    }
    constexpr friend std::ostream& operator << (std::ostream& os, const ModInt& other) {
        return os << other.x;
    }

    constexpr ModInt pow(ll k) const {
        ModInt ans = 1, tmp = x;
        while (k) {
            if (k & 1) ans *= tmp;
            tmp *= tmp;
            k >>= 1;
        }
        return ans;
    }

    constexpr ModInt inv() const {
        if (x < 1000111) {
            _precalc(1000111);
            return invs[x];
        }
        int a = x, b = MD, ax = 1, bx = 0;
        while (b) {
            int q = a/b, t = a%b;
            a = b; b = t;
            t = ax - bx*q;
            ax = bx; bx = t;
        }
        assert(a == 1);
        if (ax < 0) ax += MD;
        return ax;
    }

    static std::vector<ModInt> factorials, inv_factorials, invs;
    constexpr static void _precalc(int n) {
        if (factorials.empty()) {
            factorials = {1};
            inv_factorials = {1};
            invs = {0};
        }
        if (n > MD) n = MD;
        int old_sz = factorials.size();
        if (n <= old_sz) return;

        factorials.resize(n);
        inv_factorials.resize(n);
        invs.resize(n);

        for (int i = old_sz; i < n; ++i) factorials[i] = factorials[i-1] * i;
        inv_factorials[n-1] = factorials.back().pow(MD - 2);
        for (int i = n - 2; i >= old_sz; --i) inv_factorials[i] = inv_factorials[i+1] * (i+1);
        for (int i = n-1; i >= old_sz; --i) invs[i] = inv_factorials[i] * factorials[i-1];
    }

    static int get_primitive_root() {
        static int primitive_root = 0;
        if (!primitive_root) {
            primitive_root = [&]() {
                std::set<int> fac;
                int v = MD - 1;
                for (ll i = 2; i * i <= v; i++)
                    while (v % i == 0) fac.insert(i), v /= i;
                if (v > 1) fac.insert(v);
                for (int g = 1; g < MD; g++) {
                    bool ok = true;
                    for (auto i : fac)
                        if (ModInt(g).pow((MD - 1) / i) == 1) {
                            ok = false;
                            break;
                        }
                    if (ok) return g;
                }
                return -1;
            }();
        }
        return primitive_root;
    }

    static ModInt C(int n, int k) {
        _precalc(n + 1);
        return factorials[n] * inv_factorials[k] * inv_factorials[n-k];
    }
    
private:
    // Internal, DO NOT USE.
    // val must be in [0, 2*MD)
    constexpr inline __attribute__((always_inline)) ModInt& _set(ll v) {
        x = v >= MD ? v - MD : v;
        return *this;
    }
};
template <int MD> std::vector<ModInt<MD>> ModInt<MD>::factorials = {1};
template <int MD> std::vector<ModInt<MD>> ModInt<MD>::inv_factorials = {1};
template <int MD> std::vector<ModInt<MD>> ModInt<MD>::invs = {0};
// }}}
#line 1 "Math/Polynomial/NTT.h"
// NTT {{{
//
// Faster than NTT_chemthan.h
//
// Usage:
// auto c = multiply(a, b);
// where a and b are vector<ModInt<ANY_MOD>>
// (If mod is NOT NTT_PRIMES, it does 3 NTT and combine result)

constexpr int NTT_PRIMES[] = {998244353, 167772161, 469762049};

// assumptions:
// - |a| is power of 2
// - mint::mod() is a valid NTT primes (2^k * m + 1)
template<typename mint> void ntt(std::vector<mint>& a, bool is_inverse) {
    int n = a.size();
    if (n == 1) return;

    static const int mod = mint::mod();
    static const mint root = mint::get_primitive_root();
    assert(__builtin_popcount(n) == 1 && (mod - 1) % n == 0);

    static std::vector<mint> w{1}, iw{1};
    for (int m = w.size(); m < n / 2; m *= 2) {
        mint dw = root.pow((mod - 1) / (4 * m));
        mint dwinv = dw.inv();
        w.resize(m * 2);
        iw.resize(m * 2);
        for (int i = 0; i < m; ++i) {
            w[m + i] = w[i] * dw;
            iw[m + i] = iw[i] * dwinv;
        }
    }

    if (!is_inverse) {
        for (int m = n; m >>= 1; ) {
            for (int s = 0, k = 0; s < n; s += 2 * m, ++k) {
                for (int i = s; i < s + m; ++i) {
                    mint x = a[i], y = a[i + m] * w[k];
                    a[i] = x + y;
                    a[i + m] = x - y;
                }
            }
        }
    } else {
        for (int m = 1; m < n; m *= 2) {
            for (int s = 0, k = 0; s < n; s += 2 * m, ++k) {
                for (int i = s; i < s + m; ++i) {
                    mint x = a[i], y = a[i + m];
                    a[i] = x + y;
                    a[i + m] = (x - y) * iw[k];
                }
            }
        }
        int n_inv = mint(n).inv().x;
        for (auto& v : a) v *= n_inv;
    }
}

template<typename mint>
std::vector<mint> ntt_multiply(int sz, std::vector<mint> a, std::vector<mint> b) {
    a.resize(sz);
    b.resize(sz);
    if (a == b) { // optimization for squaring polynomial
        ntt(a, false);
        b = a;
    } else {
        ntt(a, false);
        ntt(b, false);
    }
    for (int i = 0; i < sz; ++i) a[i] *= b[i];
    ntt(a, true);
    return a;
}

template<int MOD, typename mint>
std::vector<ModInt<MOD>> convert_mint_and_multiply(
        int sz,
        const std::vector<mint>& a,
        const std::vector<mint>& b) {
    using mint2 = ModInt<MOD>;

    std::vector<mint2> a2(a.size()), b2(b.size());
    for (size_t i = 0; i < a.size(); ++i) {
        a2[i] = mint2(a[i].x);
    }
    for (size_t i = 0; i < b.size(); ++i) {
        b2[i] = mint2(b[i].x);
    }
    return ntt_multiply(sz, a2, b2);
}

long long combine(int r0, int r1, int r2, int mod) {
    using mint2 = ModInt<NTT_PRIMES[2]>;
    static const long long m01 = 1LL * NTT_PRIMES[0] * NTT_PRIMES[1];
    static const long long m0_inv_m1 = ModInt<NTT_PRIMES[1]>(NTT_PRIMES[0]).inv().x;
    static const long long m01_inv_m2 = mint2(m01).inv().x;

    int v1 = (m0_inv_m1 * (r1 + NTT_PRIMES[1] - r0)) % NTT_PRIMES[1];
    auto v2 = (mint2(r2) - r0 - mint2(NTT_PRIMES[0]) * v1) * m01_inv_m2;
    return (r0 + 1LL * NTT_PRIMES[0] * v1 + m01 % mod * v2.x) % mod;
}

template<typename mint>
std::vector<mint> multiply(const std::vector<mint>& a, const std::vector<mint>& b) {
    if (a.empty() || b.empty()) return {};
    int sz = 1, sz_a = a.size(), sz_b = b.size();
    while (sz < sz_a + sz_b) sz <<= 1;
    if (sz <= 16) {
        std::vector<mint> res(sz_a + sz_b - 1);
        for (int i = 0; i < sz_a; ++i) {
            for (int j = 0; j < sz_b; ++j) {
                res[i + j] += a[i] * b[j];
            }
        }
        return res;
    }

    int mod = mint::mod();
    std::vector<mint> res;
    if (std::find(std::begin(NTT_PRIMES), std::end(NTT_PRIMES), mod) != std::end(NTT_PRIMES)) {
        res = ntt_multiply(sz, a, b);
    } else {
        auto c0 = convert_mint_and_multiply<NTT_PRIMES[0], mint> (sz, a, b);
        auto c1 = convert_mint_and_multiply<NTT_PRIMES[1], mint> (sz, a, b);
        auto c2 = convert_mint_and_multiply<NTT_PRIMES[2], mint> (sz, a, b);

        res.resize(sz);
        for (int i = 0; i < sz; ++i) {
            res[i] = combine(c0[i].x, c1[i].x, c2[i].x, mod);
        }
    }

    res.resize(sz_a + sz_b - 1);
    return res;
}
// }}}
#line 8 "Math/tests/ntt_any_mod.test.cpp"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)
using mint = ModInt<1'000'000'007>;

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n, m; cin >> n >> m;
    vector<mint> a(n); REP(i,n) cin >> a[i];
    vector<mint> b(m); REP(i,m) cin >> b[i];

    auto c = multiply(a, b);
    for (const auto& val : c) cout << val << ' ';
    cout << endl;
    return 0;
}

Back to top page