Skip to content

Convolution with custom mod #190

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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions