naoya_t@hatenablog

いわゆるチラシノウラであります

DFT（離散フーリエ変換）とNTT（数論変換：整数（剰余環）を用いたDFT）

Algorithms Mathematics Programming Contest

先日ARC埋めの一環でATC001のC問題（高速フーリエ変換）を解いたのだけれど、こういう畳み込み問題はDFT（discrete Fourier transform；離散フーリエ変換）の代わりに NTT（number-theoretical transform；数論変換）という "整数（剰余環）を用いたDFT" でも解けるらしい。

そこで、math314氏の「 任意modでの畳み込み演算をO(n log(n))で - math314のブログ」の解説と実装例*1を追いながら自分で実装して試してみた。こないだcomplex<double>で自分で書いてみたやつ（→ATC001）とほぼ同じことをしているように見える。これは面白い。

（ここでは解説はしないので、詳しくはmath314氏の元記事を参照して頂ければと思います）

// 後述の2b.相当のコード
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

#define rep(var,n)  for(int var=0;var<(n);++var)
#define rep1(var,n)  for(int var=1;var<=(n);++var)

template<class T>
T extgcd(T a, T b, T& x, T& y) {
    for (T u=y=1,v=x=0; a; ) {
        T q = b / a;
        swap(x -= q*u, u);
        swap(y -= q*v, v);
        swap(b -= q*a, a);
    }
    return b;
}

inline long long mod_pow(long long x, long long e, long long mod) {
    long long v = 1LL;
    for ( ; e; e>>=1) {
        if (e & 1) v = (v * x) % mod;
        x = (x * x) % mod;
    }
    return v;
}

inline long long mod_inv(long long a, long long mod) {
    // return mod_pow(a, mod-2, mod); // slower
    long long x, y;
    extgcd(a, mod, x, y);
    return (mod + x % mod) % mod;
}

long long _garner(vector<long long>& xs, vector<long long>& mods) {
    int M = xs.size();

    vector<long long> coeffs(M, 1), constants(M, 0);

    for (int i=0; i<M-1; ++i) {
        long long mod_i = mods[i];
        // coffs[i] * v + constants[i] == mr[i].val (mod mr[i].first) を解く
        long long v = (xs[i] - constants[i] + mod_i) % mod_i;
        v = (v * mod_inv(coeffs[i], mod_i)) % mod_i;

        for (int j=i+1; j<M; j++) {
            long long mod_j = mods[j];
            constants[j] = (constants[j] + coeffs[j] * v) % mod_j;
            coeffs[j] = (coeffs[j] * mod_i) % mod_j;
        }
    }

    return constants.back();
}

template<typename T>
inline void bit_reverse(vector<T>& a) {
    int n = a.size();
    int i = 0;
    for (int j=1; j<n-1; ++j) {
        for (int k = n >> 1; k >(i ^= k); k >>= 1);
        if (j < i) swap(a[i], a[j]);
    }
}


template<long long mod, long long primitive_root>
class NTT {
public:
    long long get_mod() { return mod; }

    void _ntt(vector<long long>& a, int sign) {
        const int n = a.size();
        assert((n ^ (n&-n)) == 0); //n = 2^k

        const long long g = primitive_root; // g is primitive root of mod

        long long tmp = (mod - 1) * mod_inv(n, mod) % mod; // -1/n
        long long h = mod_pow(g, tmp, mod); // ^n√g
        if (sign == -1) h = mod_inv(h, mod);

        bit_reverse(a);

        for (int m=1; m<n; m<<=1) {
            const int m2 = 2 * m;
            long long _base = mod_pow(h, n/m2, mod);
            long long _w = 1;
            for (int x=0; x<m; ++x) {
                for (int s=x; s<n; s+=m2) {
                    long long u = a[s];
                    long long d = (a[s + m] * _w) % mod;
                    a[s] = (u+d) % mod;
                    a[s+m] = (u-d+mod) % mod;
                }
                _w = (_w * _base) % mod;
            }
        }
    }
    void ntt(vector<long long>& input) {
        _ntt(input, 1);
    }
    void intt(vector<long long>& input) {
        _ntt(input, -1);

        const long long n_inv = mod_inv(input.size(), mod);
        for (auto &x : input) x = (x * n_inv) % mod;
    }

    // 畳み込み演算を行う
    vector<long long> convolution(const vector<long long>& a, const vector<long long>& b){
        int result_size = a.size() + b.size() - 1;
        int n = 1; while (n < result_size) n <<= 1;

        vector<long long> _a = a, _b = b;
        _a.resize(n, 0);
        _b.resize(n, 0);

        ntt(_a);
        ntt(_b);
        for (int i=0; i<n; ++i) _a[i] = (_a[i] * _b[i]) % mod;
        intt(_a);

        _a.resize(result_size);
        return _a;
    }
};


vector<long long> convolution_ntt(vector<long long>& a, vector<long long>& b, long long mod=1224736769LL) {
    for (auto &x : a) x %= mod;
    for (auto &x : b) x %= mod;

    ll maxval = max(a.size(), b.size()) * *max_element(a.begin(), a.end()) * *max_element(b.begin(), b.end());
    if (maxval < 1224736769) {
    	NTT<1224736769, 3> ntt3;
        return ntt3.convolution(a, b);
    }

    NTT<167772161, 3> ntt1;
    NTT<469762049, 3> ntt2;
    NTT<1224736769, 3> ntt3;

    vector<long long> x1 = ntt1.convolution(a, b);
    vector<long long> x2 = ntt2.convolution(a, b);
    vector<long long> x3 = ntt3.convolution(a, b);

    vector<long long> ret(x1.size());
    vector<long long> mods { 167772161, 469762049, 1224736769, mod };
    for (int i=0; i<x1.size(); ++i) {
        vector<long long> xs { x1[i], x2[i], x3[i], 0 };
        ret[i] = _garner(xs, mods);
    }

    return ret;
}


int main() {
    int N; cin >> N; // 1-100000

    vector<ll> g(N+1), h(N+1);

    g[0] = h[0] = 0;
    rep(i,N){
        int A,B; cin >> A >> B;
        g[1+i] = A;
        h[1+i] = B;
    }

    vector<ll> f = convolution_ntt(g, h, 1000000007LL);
    rep1(k, N*2){
        cout << f[k] << endl;
    }

    return 0;
}

1. math314氏の実装例を適宜修正したもの（3つのNTTの結果をGarner法でまとめる）。

NTTを3回、それぞれ別々の法 (167772161, 469762049, 1224736769) を用いて実行する。
この3つの素数のチョイスは 167772161 \cdot 469762049 \cdot 1224736769 > 2^{86} > 2^{84} \simeq \max(size_a, size_b) \cdot int\_max^2 ということらしい。（\max(size_a, size_b)=2^{22}\simeq4\cdot10^6としている。INT_MAXまでの非負整数であれば復元可能。ただし64bitを超えるのでlong longに収まりきれず中国剰余定理では無理で、代わりにGarner法が用いられる。）

◎math314氏のコードの修正点：
  1. const int g = 3; は g = primitive_root; とするか gの代わりに primitive_root をそのまま使った方が良いかと。ここで用いられる3つの素数は全て原始根（primitive root）として3を持つので動作には支障ないのだけれど*2。
  2. その次の行でhを求める際、(mod - 1) / n の部分は (mod - 1) * mod_inv(n, mod) % mod じゃないといけない気がする。
  • 1a. 演算に剰余環演算クラスを使用

　→AC (956ms) https://atc001.contest.atcoder.jp/submissions/2628347

  • 1b. 剰余環演算クラスを使わずにlong long演算で

　→AC (758ms) https://atc001.contest.atcoder.jp/submissions/2628636

2. この例題は各項が最大でも1e9までの値にしかならないのでNTT<1224736769,3>の1つだけ行う。

  • 2a. 演算に剰余環演算クラスを使用。

　→AC (454ms) https://atc001.contest.atcoder.jp/submissions/2628363

  • 2b. 剰余環演算クラスを使わずにlong long演算で

　→AC (427ms) https://atc001.contest.atcoder.jp/submissions/2628643

3 (参考). 先日のFFT

  • 先日の実装

　→AC (866ms) https://atc001.contest.atcoder.jp/submissions/2583908

  • 再帰ではなくbit_reverseを使って書き直したもの

　→AC (441ms) https://atc001.contest.atcoder.jp/submissions/2628595

その他

剰余環における逆元を求める mod_inv() の計算速度

普段はフェルマーの小定理を用いる方法（(p-2)乗するやつ）を使っているコピペしているのだけれど、拡張ユークリッド互除法を用いた計算のほうが2倍ぐらい速い。

typedef long long ll;

ll mod_pow(ll x, ll e, ll p) {
    // x^e mod pを求める
    ll v = 1LL;
    for ( ; e; e>>=1) {
        if (e & 1) v = (v * x) % p;
        x = (x * x) % p;
    }
    return v;
}

ll mod_inv_fermat(ll a, ll p) {
    // フェルマーの小定理を用いた逆元計算
    return mod_pow(a, p-2, p);  // a^(p-2) mod p
}
typedef long long ll;

ll extgcd(ll a, ll b, ll& x, ll& y) {
    // 拡張ユークリッド互除法
    // ax+by=gcd(a,b) を満たすx,yを1組求めたい時に
    for (ll u=y=1,v=x=0; a; ) {
        ll q = b / a;
        swap(x -= q*u, u);
        swap(y -= q*v, v);
        swap(b -= q*a, a);
    }
    return b;
}

ll mod_inv_extgcd(ll a, ll mod) {
    // 拡張ユークリッド互除法を用いた逆元計算
    ll x, y;
    extgcd(a, mod, x, y);
    return (mod + x % mod) % mod;
}

求めたい値は 1/a（あるいはこれにmodの整数倍を足したもの）なので、
\begin{align}
x &= \frac{1}{a} + kM \\
x - kM &= \frac{1}{a} \\
a(x - kM) &= 1 \\
ax + M(-ak) &= 1 \\
y=-akとおいて \\
ax + My &= 1 = gcd(a,M)
\end{align}
↑aとMは互いに素（Mは素数だし0\lt a\lt Mだし）なのでgcd(a,M)=1である。
ということで、ax+My=1を満たす(x,y)を1つ求める拡張ユークリッド互除法、すなわちmod_inv_extgcd(a, M) を用いて x を求め（yはk部分だから捨てていい）、0\le x\lt M の範囲になるように補正すれば良い。

参考資料

追記

けんちょん先生がちょうど逆元計算の説明をしてくれてる

*1:動かすには何箇所か修正が必要

*2:ちなみに1000000007の原始根はWolfram Alpha先生によると5である。