任意模数NTT学习小记

问题

求两个多项式$A(x)$和$B(x)$对一个不是NTT模数的数取模的结果。

拆系数FFT

设置一个阈值$W$（通常设置为$2^{15}$），将$A(x),B(x)$拆分成$A=aW+b,B=cW+d$，其中$a,b,c,d$均为多项式。则$$AB=ac{W}^2+(ad+bc)W+bd$$

需要做$7$次DFT。

将FFT中DFT的次数优化为两次

对于两个实系数多项式$A(x),B(x)$，现在要计算它们的乘积。设$$F(x)=A(x)+iB(x),G(x)=A(x)-iB(x)$$

令$\omega$表示$n$次单位根，$F_{DFT}(k)=F({\omega }^k),G_{DFT}(k)=G({\omega }^k)$。则显然有$$F_{DFT}(k)=A({\omega }^k)+iB({\omega }^k)=\sum _{j=0}^{n-1}(a_j+i{b}_j){\omega }^{jk}$$

考虑如何通过$F$的点值得到$G$的点值。令$conj(w)$表示$w$的共轭，那么$$\begin{array}{rl}{G}_{DFT}(k)& =\sum _{j=0}^{n-1}(a_j-i{b}_j){\omega }^{jk}\\ & =\sum _{j=0}^{n-1}(a_j-i{b}_j)(\cos \frac{2\pi jk}{n}+i\sin \frac{2\pi jk}{n})\\ & =\sum _{j=0}^{n-1}(a_j\cos \frac{2\pi jk}{n}+b_j\sin \frac{2\pi jk}{n})+i(a_j\sin \frac{2\pi jk}{n}-b_j\cos \frac{2\pi jk}{n})\\ & =conj\left(\sum _{j=0}^{n-1}\right(a_j\cos \frac{-2\pi jk}{n}-b_j\sin \frac{-2\pi jk}{n})+i(a_j\sin \frac{-2\pi jk}{n}+b_j\cos \frac{-2\pi jk}{n}\left)\right)\\ & =conj(\sum _{j=0}^{n-1}(a_j+i{b}_j)(\cos \frac{-2\pi jk}{n}+i\sin \frac{-2\pi jk}{n}\left)\right)\\ & =conj(\sum _{j=0}^{n-1}(a_j+i{b}_j){\omega }^{-jk})\\ & =conj(\sum _{j=0}^{n-1}(a_j+i{b}_j){\omega }^{j(n-k)})\\ & =conj(F_{DFT}(n-k))\end{array}$$

把$F,G$的点值求出来，就可以解出$A,B$的点值。从而把FFT中DFT的次数从三次优化为两次。

对拆系数FFT的优化

从上面的推导容易发现，若知道一个多项式的点值，则可以求出其共轭多项式的点值。将这个优化用于拆系数FFT。

注意到$$(a+bi)(c+di)=(ac-bd)+i(ad+bc)$$

$$(a-bi)(c+di)=(ac+bd)+i(ad-bc)$$

可以先求出$(a+bi)$和$(c+di)$的点值。由于$(a+bi)$和$(a-bi)$互为共轭，可以通过$(a+bi)$的点值推出$(a-bi)$的点值。求出上面两式后再IDFT回去，就可以得到$ac,bd,ad+bc$的值了。这样总共需要做四次DFT。

代码

#include<bits/stdc++.h>
using namespace std;

typedef long long LL;

const int W = 32768;
const int N = 400005;
const long double pi = acos(-1);

int n, m, p, rev[N], L, ans[N];

struct com
{
    long double x,y;
    com operator + (const com & d) const {return (com){x + d.x, y + d.y};}
    com operator - (const com & d) const {return (com){x - d.x, y - d.y};}
    com operator * (const com & d) const {return (com){x * d.x - y * d.y, x * d.y + y * d.x};}
    com operator / (const long double & d) const {return (com){x / d, y / d};}
    com operator ~ () const {return (com){x, -y};}
}A[N], B[N], C[N];

void pre()
{
	int lg = 0;
	for (L = 1; L <= n + m; L <<= 1, lg++);
	for (int i = 0; i < L; i++) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (lg - 1));
}

LL num(long double x)
{
	return x < 0 ? (LL)(x - 0.5) : (LL)(x + 0.5);
}

void fft(com * a, int f)
{
	for (int i = 0; i < L; i++) if (i < rev[i]) swap(a[i], a[rev[i]]);
	for (int i = 1; i < L; i <<= 1)
	{
		com wn = (com){cos(pi / i), f * sin(pi / i)};
		for (int j = 0; j < L; j += (i << 1))
		{
			com w = (com){1, 0};
			for (int k = 0; k < i; k++)
			{
				com u = a[j + k], v = a[j + k + i] * w;
				a[j + k] = u + v; a[j + k + i] = u - v;
				w = w * wn;
			}
		}
	}
	if (f == -1) for (int i = 0; i < L; i++) a[i] = a[i] / L;
}

int main()
{
	scanf("%d%d%d", &n, &m, &p);
	for (int i = 0; i <= n; i++)
	{
		int x; scanf("%d", &x);
		A[i] = (com){x / W, x % W};
	}
	for (int i = 0; i <= m; i++)
	{
		int x; scanf("%d", &x);
		C[i] = (com){x / W, x % W};
	}
	pre();
	fft(A, 1); fft(C, 1);
	for (int i = 0; i < L; i++) B[i] = ~A[(L - i) % L];
	for (int i = 0; i < L; i++) A[i] = A[i] * C[i], B[i] = B[i] * C[i];
	fft(A, -1); fft(B, -1);
	for (int i = 0; i <= n + m; i++)
	{
		LL x1 = num(A[i].x), x2 = num(B[i].x), y1 = num(A[i].y), y2 = num(B[i].y);
		LL ac = (x1 + x2) / 2, bd = x2 - ac, bcad = y1;
		ans[i] = (ac % p * W % p * W % p + bcad * W % p + bd % p) % p;
	}
	for (int i = 0; i <= n + m; i++) printf("%d ", ans[i]);
	return 0;
}