任意mod畳み込み (lib/fft/ntt_mod.hpp)

• View this file on GitHub
• Last update: 2025-11-18 14:07:26+09:00
• Include: #include "lib/fft/ntt_mod.hpp"
• Link: https://asako.growi.cloud/compro/NTT
• Link: https://math314.hateblo.jp/entry/2015/05/07/014908

Depends on

• 畳み込み (lib/fft/ntt.hpp)
• lib/internal/internal_fft.hpp
• Montgomery modular multiplication (lib/internal/internal_math.hpp)
• lib/internal/internal_type_traits.hpp
• lib/math/modint.hpp
• lib/template/template.hpp

Verified with

• test/yosupo/convolution/convolution_mod.test.cpp

Code

#pragma once
#include "fft/ntt.hpp"
#include "template/template.hpp"

/**
 * @brief 任意mod畳み込み
 *
 * @tparam mod
 * @tparam T
 * @param a
 * @param b
 * @return std::vector<T>
 *
 * @see https://math314.hateblo.jp/entry/2015/05/07/014908
 * @see https://asako.growi.cloud/compro/NTT
 */
template <unsigned int mod = 998244353, class T, std::enable_if_t<std::is_integral<T>::value> * = nullptr>
std::vector<T> convolution_mod(const std::vector<T> &a, const std::vector<T> &b) {
    int n = int(a.size()), m = int(b.size());
    if (!n || !m) return {};

    static constexpr std::uint64_t MOD1 = 754974721;  // 2^24
    static constexpr std::uint64_t MOD2 = 167772161;  // 2^25
    static constexpr std::uint64_t MOD3 = 469762049;  // 2^26
    static constexpr std::uint64_t M2M3 = MOD2 * MOD3;
    static constexpr std::uint64_t M1M3 = MOD1 * MOD3;
    static constexpr std::uint64_t M1M2 = MOD1 * MOD2;
    static constexpr __uint128_t M1M2M3 = (__uint128_t)MOD1 * MOD2 * MOD3;

    static constexpr std::uint64_t i1 = internal::inv_gcd(MOD2 * MOD3, MOD1).second;
    static constexpr std::uint64_t i2 = internal::inv_gcd(MOD1 * MOD3, MOD2).second;
    static constexpr std::uint64_t 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<T> c(n + m - 1);
    for (int i = 0; i < n + m - 1; ++i) {
        __uint128_t x = (c1[i] * i1) % MOD1 * (__uint128_t)M2M3;
        x += (c2[i] * i2) % MOD2 * (__uint128_t)M1M3;
        if (x >= M1M2M3) x -= M1M2M3;
        x += (c3[i] * i3) % MOD3 * (__uint128_t)M1M2;
        if (x >= M1M2M3) x -= M1M2M3;
        c[i] = x % mod;
    }

    return c;
}

Traceback (most recent call last):
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/resolver.py", line 291, in resolve
    bundled_code = language.bundle(path, basedir=basedir)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus.py", line 242, in bundle
    bundler.update(path)
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus_bundle.py", line 479, in update
    self._resolve(pathlib.Path(included), included_from=path)
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus_bundle.py", line 286, in _resolve
    raise BundleErrorAt(path, -1, "no such header")
competitive_verifier.oj.verify.languages.cplusplus_bundle.BundleErrorAt: fft/ntt.hpp: line -1: no such header

Back to top page