FFT优化技巧——MTT

2018-08-05 18:35:43


一、MTT的适用范围

模板题:P4245 【模板】任意模数NTT

现在算法竞赛的毒瘤们常常要求对某一个数取膜。如果使用FFT直接求解,long double也会出现不小的误差。如果取膜的数是一个比DFT长度大的、形如$a\times2^k+1$的数,那么我们可以求出原根,代替复根,进行DFT和IDFT。这种算法被称为数论变换(Number Theorem Transform(NTT))。

然而,这样的膜数毕竟很少,NTT的适用面较窄。如果遇到随便的一个膜数,一种方法是选择三个满足NTT性质的膜数,分别求出这三个膜数意义下的卷积,再使用CRT确定答案。

这种方法的效率并不让人满意。同时,对于初学DFT的选手们来说,原根、CRT这样的知识也难以迅速全部理解。其实,要对所得的多项式取膜并保证精度,还有一种更易懂,更易实现,效率比三膜数NTT更高的实现方法。这种方法就是MTT(毛啸变换(Matthew Transform))。取膜和不取膜的多项式卷积都可以使用MTT。

二、MTT的基本思想

拆系数。

以洛谷P4245任意膜数NTT为例,发现系数的范围在$10^9$以内。可将两个多项式的系数都拆成$a\times32768+b$的形式。

比如计算多项式$A(x)$和$B(x)$在$mod\space p$意义下的卷积。将$A(x)$的系数序列$a$拆为两个序列$k[a]$和$b[a]$,其中$k[a_i]=a_i\div32768,b[a_i]=a_i\space mod\space32768$。设这两个新序列形成的多项式为$C,D$。对$B$也进行相同操作,形成多项式$E,F$。

不难发现,整个多项式在$n$处的点值表示即为$$A(n)\cdot B(n)=(32768C(n)+D(n))\cdot(32768E(n)+F(n))$$$$=1073741824C(n)\cdot E(n)+32768(C(n)\cdot F(n)+E(n)\cdot D(n))+D(n)\cdot F(n)$$

为了保证精度,我们计算$C(n)\cdot E(n),C(n)\cdot F(n),E(n)\cdot D(n),D(n)\cdot F(n)$四个卷积。再按上式计算点值。

三、MTT的实现细节

我们预先算出上面$C,D,E,F$多项式的DFT,再求出4个卷积的点值表示,再进行4次IDFT算出四个卷积,然后按贡献相加。这种最朴素的MTT共需要4次复数意义下的DFT和4次复数意义下的IDFT。

注意到$C(n)\cdot F(n),E(n)\cdot D(n)$的贡献相同,可以直接计算$C(n)\cdot F(n)+E(n)\cdot D(n)$的点值表示,再经过一次IDFT求出其系数。这样可以减少一次IDFT。

既然是复数计算,就要考虑精度。这里一定要使用$std::sin$和$std::cos$,一定要加$std::$。这样精度更高。

#include<cstdio>
#include<cmath>
#include<cctype>
#include<cstring>
#include<algorithm>
#define LL long long

int Read()
{
    int x=0;char c=getchar();
    while(!isdigit(c))
    {
        c=getchar();
    }
    while(isdigit(c))
    {
        x=x*10+(c^48);
        c=getchar();
    }
    return x;
}

const long double PI=acos(-1);

struct Complex
{
    long double r,i;
    Complex(long double R=0,long double I=0)
    {
        r=R;
        i=I;
    }
}ak[1000005],ab[1000005],bk[1000005],bb[1000005],A[1000005],B[1000005],C[1000005],D[1000005];

Complex operator * (Complex a,Complex b)
{
    return Complex(a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r);
}

Complex operator + (Complex a,Complex b)
{
    return Complex(a.r+b.r,a.i+b.i);
}

Complex operator - (Complex a,Complex b)
{
    return Complex(a.r-b.r,a.i-b.i);
}

int rev[1000005];
void getrev(int bit)
{
    for(int i=0;i<(1<<bit);i++)
        rev[i]=(rev[i>>1]>>1)|((i&1)<<(bit-1));
}

void fft(Complex* a,int n,int dft)
{
    for(int i=0;i<n;i++)
        if(i<rev[i])
            std::swap(a[i],a[rev[i]]);
    for(int step=1;step<n;step<<=1)
    {
        Complex omega(std::cos(dft*PI/step),std::sin(dft*PI/step));
        for(int j=0;j<n;j+=step<<1)
        {
            Complex omega_k(1,0);
            for(int k=j;k<j+step;k++)
            {
                Complex x=a[k];
                Complex y=omega_k*a[k+step];
                a[k]=x+y;
                a[k+step]=x-y;
                omega_k=omega_k*omega;
            }
        }
    }
    if(dft==-1)
    {
        for(int i=0;i<n;i++)
            a[i].r/=n;
    }
}

int num_a[100005],num_b[100005];
int main()
{
    int n=Read(),m=Read(),p=Read();
    int bit=1,s=2;
    for(bit=1;(1<<bit)<n+m+1;++bit)
        s<<=1;
    for(int i=0;i<=n;++i)
    {
        num_a[i]=Read()%p;
        ak[i].r=num_a[i]>>15;
        ab[i].r=num_a[i]&0x7fff;
    }
    for(int i=0;i<=m;++i)
    {
        num_b[i]=Read()%p;
        bk[i].r=num_b[i]>>15;
        bb[i].r=num_b[i]&0x7fff;
    }
    getrev(bit);
    fft(ak,s,1);
    fft(ab,s,1);
    fft(bk,s,1);
    fft(bb,s,1);
    for(int i=0;i<s;++i)
    {
        A[i]=ak[i]*bk[i];
        B[i]=ak[i]*bb[i]+ab[i]*bk[i];
        D[i]=ab[i]*bb[i];
    }
    fft(A,s,-1);
    fft(B,s,-1);
    fft(D,s,-1);
    for(int i=0;i<=m+n;++i)
    {
        printf("%lld ",((((LL)(A[i].r/1+0.5)%p)<<30)+(((LL)(B[i].r/1+0.5)%p)<<15)+(LL)(D[i].r/1+0.5)%p)%p);
    }
    return 0;
}

四、进一步的优化

事实上,MTT还可以继续优化,在毛啸的论文中,最多可以减至1.5次DFT和2次IDFT。

这里让我们来一步步优化。

首先介绍一种将DFT两两合并的技巧,

考虑对长度为$n=2^n$的实多项式$A(x),B(x)$进行DFT。

定义:$$P(x)=A(x)+iB(x)$$$$Q(x)=A(x)-iB(x)$$

其中$i=\sqrt{-1}$。

那么有:$$DFT(A)=\frac{DFT(P)+DFT(Q)}{2}$$$$DFT(B)=i\frac{DFT(P)-DFT(Q)}{2}$$

如果能在求出$DFT(P(x))$后迅速求出$DFT(Q(x))$就好了,这样两次DFT可以优化为一次。

我们可以进行这样的推导:

即得易见平凡,仿照上例显然。留作习题答案略,读者自证不难。
反之亦然同理,推论自然成立。略去过程Q.E.D.,由上可知证毕。

(以后会完善证明过程)

可知:$$Q(\omega^k_n)=conj(P(\omega^{n-k}_n))$$

其中$conj(x)$表示$x$的共轭复数。这样,可以把两次实多项式的DFT合并成一次DFT。IDFT也可以作同样的优化。

这样,4次DFT可以优化成2次,3次IDFT中选2次合并成1次,总共2次IDFT。这样的效率已经相当高了。

这两次DFT还可以继续优化。

(鸽子:别急,我正在写呢)

五、神仙代码

毛啸在UOJ多项式乘法一题中的提交记录

#include <bits/stdc++.h>

using namespace std;

#define REP(i, a, b) for (int i = (a), _end_ = (b); i < _end_; ++i)
#define debug(...) fprintf(stderr, __VA_ARGS__)
#define mp make_pair
#define x first
#define y second
#define pb push_back
#define SZ(x) (int((x).size()))
#define ALL(x) (x).begin(), (x).end()

template<typename T> inline bool chkmin(T &a, const T &b) { return a > b ? a = b, 1 : 0; }
template<typename T> inline bool chkmax(T &a, const T &b) { return a < b ? a = b, 1 : 0; }

typedef long long LL;

const int oo = 0x3f3f3f3f;

const int Mod = 1e9 + 7;

const int max0 = 262144;

struct comp
{
    double x, y;

    comp(): x(0), y(0) { }
    comp(const double &_x, const double &_y): x(_x), y(_y) { }

};

inline comp operator+(const comp &a, const comp &b) { return comp(a.x + b.x, a.y + b.y); }
inline comp operator-(const comp &a, const comp &b) { return comp(a.x - b.x, a.y - b.y); }
inline comp operator*(const comp &a, const comp &b) { return comp(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); }
inline comp conj(const comp &a) { return comp(a.x, -a.y); }

const double PI = acos(-1);

int N, L;

comp w[max0 + 5];
int bitrev[max0 + 5];

void fft(comp *a, const int &n)
{
    REP(i, 0, n) if (i < bitrev[i]) swap(a[i], a[bitrev[i]]);
    for (int i = 2, lyc = n >> 1; i <= n; i <<= 1, lyc >>= 1)
        for (int j = 0; j < n; j += i)
        {
            comp *l = a + j, *r = a + j + (i >> 1), *p = w;
            REP(k, 0, i >> 1)
            {
                comp tmp = *r * *p;
                *r = *l - tmp, *l = *l + tmp;
                ++l, ++r, p += lyc;
            }
        }
}

inline void fft_prepare()
{
    REP(i, 0, N) bitrev[i] = bitrev[i >> 1] >> 1 | ((i & 1) << (L - 1));
    REP(i, 0, N) w[i] = comp(cos(2 * PI * i / N), sin(2 * PI * i / N));
}

inline void conv(int *x, int *y, int *z)
{
    REP(i, 0, N) (x[i] += Mod) %= Mod, (y[i] += Mod) %= Mod;
    static comp a[max0 + 5], b[max0 + 5];
    static comp dfta[max0 + 5], dftb[max0 + 5], dftc[max0 + 5], dftd[max0 + 5];

    REP(i, 0, N) a[i] = comp(x[i] & 32767, x[i] >> 15);
    REP(i, 0, N) b[i] = comp(y[i] & 32767, y[i] >> 15);
    fft(a, N), fft(b, N);
    REP(i, 0, N)
    {
        int j = (N - i) & (N - 1);
        static comp da, db, dc, dd;
        da = (a[i] + conj(a[j])) * comp(0.5, 0);
        db = (a[i] - conj(a[j])) * comp(0, -0.5);
        dc = (b[i] + conj(b[j])) * comp(0.5, 0);
        dd = (b[i] - conj(b[j])) * comp(0, -0.5);
        dfta[j] = da * dc;
        dftb[j] = da * dd;
        dftc[j] = db * dc;
        dftd[j] = db * dd;
    }
    REP(i, 0, N) a[i] = dfta[i] + dftb[i] * comp(0, 1);
    REP(i, 0, N) b[i] = dftc[i] + dftd[i] * comp(0, 1);
    fft(a, N), fft(b, N);
    REP(i, 0, N)
    {
        int da = (LL)(a[i].x / N + 0.5) % Mod;
        int db = (LL)(a[i].y / N + 0.5) % Mod;
        int dc = (LL)(b[i].x / N + 0.5) % Mod;
        int dd = (LL)(b[i].y / N + 0.5) % Mod;
        z[i] = (da + ((LL)(db + dc) << 15) + ((LL)dd << 30)) % Mod;
    }
}

int main()
{
#ifndef ONLINE_JUDGE
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
#endif
    int n, m;
    static int a[max0 + 5], b[max0 + 5], c[max0 + 5];
    scanf("%d%d", &n, &m), ++n, ++m;
    REP(i, 0, n) scanf("%d", a + i);
    REP(i, 0, m) scanf("%d", b + i);
    L = 0;
    for ( ; (1 << L) < n + m - 1; ++L);
    N = 1 << L;
    fft_prepare();
    conv(a, b, c);
    REP(i, 0, n + m - 1) (c[i] += Mod) %= Mod, printf("%d ", c[i]);
    printf("\n");
    return 0;
}