Skip to content

Convolution with custom mod #190

@anhkha2003

Description

@anhkha2003

Sometimes we want to do convolution modulo 10^9 + 7, e.g. https://judge.yosupo.jp/problem/convolution_mod_1000000007

convolution_ll will overflow. Here is my adjusted code that works for a custom modulus. Can we add a similar function to the library?

const long long MOD = 1e9 + 7;

vector<long long> multiply(vector<long long>& a, vector<long long>& b) {
    int n = int(a.size()), m = int(b.size());
    if (!n || !m) return {};

    static constexpr unsigned long long MOD1 = 754974721;  // 2^24
    static constexpr unsigned long long MOD2 = 167772161;  // 2^25
    static constexpr unsigned long long MOD3 = 469762049;  // 2^26
    static constexpr __int128 M2M3 = MOD2 * MOD3;
    static constexpr __int128 M1M3 = MOD1 * MOD3;
    static constexpr __int128 M1M2 = MOD1 * MOD2;
    static constexpr __int128 M1M2M3 = (__int128)MOD1 * MOD2 * MOD3;

    static constexpr long long i1 =
        internal::inv_gcd(MOD2 * MOD3, MOD1).second;
    static constexpr long long i2 =
        internal::inv_gcd(MOD1 * MOD3, MOD2).second;
    static constexpr long long i3 =
        internal::inv_gcd(MOD1 * MOD2, MOD3).second;
        
    static constexpr int MAX_AB_BIT = 24;
    static_assert(MOD1 % (1ull << MAX_AB_BIT) == 1, "MOD1 isn't enough to support an array length of 2^24.");
    static_assert(MOD2 % (1ull << MAX_AB_BIT) == 1, "MOD2 isn't enough to support an array length of 2^24.");
    static_assert(MOD3 % (1ull << MAX_AB_BIT) == 1, "MOD3 isn't enough to support an array length of 2^24.");
    assert(n + m - 1 <= (1 << MAX_AB_BIT));

    auto c1 = convolution<MOD1>(a, b);
    auto c2 = convolution<MOD2>(a, b);
    auto c3 = convolution<MOD3>(a, b);

    std::vector<long long> c(n + m - 1);
    for (int i = 0; i < n + m - 1; i++) {
        __int128 x = 0;
        x += (c1[i] * i1) % MOD1 * M2M3;
        x += (c2[i] * i2) % MOD2 * M1M3;
        x += (c3[i] * i3) % MOD3 * M1M2;
        x %= M1M2M3;

        x %= MOD;
        c[i] = x;
    }

    return c;
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions