-
Notifications
You must be signed in to change notification settings - Fork 261
Open
Description
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
Labels
No labels