离散傅里叶变换

实在不知道标题叫啥了,总不能就 DFT 三个字母吧。是最近把 \(\rm DFT\) 的过程自己推了一遍后的浅显理解,以及更快的 \(\rm FFT\) 板子。因为转置原理对于实现的优化实在是太多了,所以我也简单学了一下,不过真的啥都不学不懂。

DFT 在干啥

首先我们要知道 \(\rm DFT\) 是什么,简单来说,它是一组复数 \(\left<a_{0..n-1}\right>\) 到另一组复数 \(\left<b_{0..n-1}\right>\) 的线性变换。式子长这样,其中 \(i\) 是虚数单位:

\[b_j=\sum_{k=0}^{n-1} a_ke^{-\frac{i2\pi jk}{n}} \]

这个 \(e^{-\frac{i2\pi jk}{n}}\) 其实就是 \(n\) 次单位根的 \(jk\) 次方,\(\omega_n^{jk}\)。可以写出它对应的矩阵:

\[\begin{bmatrix} \omega_n^{0\times 0}&\omega_n^{0\times 1}&\omega_n^{0\times 2}&\cdots&\omega_n^{0\times (n-1)}\\ \omega_n^{1\times 0}&\omega_n^{1\times 1}&\omega_n^{1\times 2}&\cdots&\omega_n^{1\times (n-1)}\\ \vdots&\vdots&\vdots&\ddots&\vdots\\\omega_n^{(n-1)\times 0}&\omega_n^{(n-1)\times 1}&\omega_n^{(n-1)\times 2}&\cdots&\omega_n^{(n-1)\times (n-1)} \end{bmatrix}\]

也叫范德蒙德矩阵。这个矩阵可以在后面的转置原理里用到,这里暂且按下不表。

接下来就来看看它在干什么吧。考虑两个序列 \(\left<f_{0..n-1}\right>,\left<g_{0..n-1}\right>\) 进行 \(\rm DFT\) 后进行点乘,本文约定序列 \(\left<f_{0..n-1}\right>\) 进行 \(\rm DFT\) 后得到的序列为 \(\left<\hat{f}_{0..n-1}\right>\)

\[\begin{aligned} \hat{f}_k\times \hat{g}_k&=\left(\sum_{i=0}^{n-1} f_i\omega_n^{ik}\right)\left(\sum_{j=0}^{n-1}g_j\omega_n^{jk}\right)\\ &=\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j\omega_n^{(i+j)k}\\ &=\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j\omega_n^{(i+j)k\bmod n}\\ &=\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j\sum_{p=0}^{n-1}[(i+j)\bmod n=p]\omega_n^{pk}\\ &=\sum_{p=0}^{n-1}\omega_n^{pk}\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j[(i+j)\bmod n=p] \end{aligned}\]

发现也是 \(\rm DFT\) 的形式,考虑设 \(\hat{h}_k=\hat{f}_k\times \hat{g}_k\),则可以发现:

\[h_k=\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j[(i+j)\bmod n=k] \]

也就是 \(f_i,g_j\) 在模 \(n\) 意义下的循环卷积。换句话说,我们把 \(\hat{f},\hat{g}\) 点乘就得到了 \(f,g\) 循环卷积后的序列 \(\rm DFT\) 后的结果。如果我们能找到一个逆 \(\rm DFT\),或者说 \(\rm IDFT\) 的过程,就可以实现循环卷积的计算了。考虑从 \(h_k\) 的表达式入手:

\[\begin{aligned} h_k&=\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j[(i+j)\bmod n=k]\\ &=\frac{1}{n}\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}f_ig_j\sum_{p=0}^{n-1}\omega_{n}^{(i+j-k)p}\\ &=\frac{1}{n}\sum_{p=0}^{n-1}\omega_n^{-kp}\left(\sum_{i=0}^{n-1}f_i\omega_n^{ip}\right)\left(\sum_{j=0}^{n-1}g_j\omega_n^{jp}\right)\\ &=\frac{1}{n}\sum_{p=0}^{n-1}\omega_n^{-kp} \hat{f}_p\times \hat{g}_p\\ &=\frac{1}{n}\sum_{p=0}^{n-1}\hat{h}_p\omega_n^{-kp}\\ &=\frac{1}{n}\sum_{p=0}^{n-1}\hat{h}_p\omega_n^{k(n-p)}\\ &=\frac{1}{n}\hat{h}_0\omega_n^{k\times 0}+\sum_{p=1}^{n-1}\hat{h}_{n-p}\omega_n^{kp} \end{aligned}\]

可以发现,\(\rm IDFT\) 的过程既可以通过把 \(\rm DFT\) 的单位根取反得到,也可以通过把序列 \(1\sim n-1\) 处的值反转得到再进行 \(\rm DFT\) 得到。

DFT 该咋做

常见的实现是 \(\rm FFT\),它的限制是要求循环卷积的模数是 \(2\) 的次幂,通过把这个模数取到大于能得到的最高次幂的方法,我们可以让循环卷积变为普通加法卷积。如果需要实现任意模数的循环卷积,我们需要再推一推式子,即 \(\rm Bluestein's\ algorithm\)。我们分别来看。

首先介绍一下 \(\rm FFT\)。它的基本思路是分治,因为 \(n\)\(2\) 的次幂,所以我们可以每次分成两半来算。考虑设:

\[f_{0,i}=f_{2i},f_{1,i}=f_{2i+1} \]

\[\begin{aligned} \hat{f}_i&=\sum_{j=0}^{n-1}f_j\omega_n^{ij}\\ &=\sum_{j=0}^{n-1}([j\bmod 2=0]f_j\omega_n^{ij}+[j\bmod 2=1]f_j\omega_n^{ij})\\ &=\sum_{j=0}^{n-1}([j\bmod 2=0]f_j\omega_{n/2}^{ij/2}+\omega_n^{i}[j\bmod 2=1]f_j\omega_{n/2}^{i(j-1)/2})\\ &=\hat{f}_{0,i}+\omega_n^i\hat{f}_{1,i} \end{aligned}\]

从而找到一个分治结构。此外,还可以发现 \(\hat{f}_i\)\(\hat{f}_{i+n/2}\) 式子有一定的相似性:

\[\hat{f}_i=\hat{f}_{0,i}+\omega_n^i\hat{f}_{1,i},\hat{f}_{i+n/2}=\hat{f}_{0,i}-\omega_n^i\hat{f}_{1,i}(i<n/2) \]

其实有 \(\hat{f}_{0/1,i+n/2}=\hat{f}_{0/1,i}(i<n/2)\)。从而可以一次计算两个的值,且每次分治长度会减半,从而我们得到一个:

\[T(n)=2T(n/2)+\mathcal{O}(n),T(n)=\mathcal{O}(n\log n) \]

的做法。

递归常数比较大,怎么办呢?模拟一下递归的过程,我们可以找到一个规律。从而可以让我们模拟递归的过程。具体来讲,如果令 \(g_i=f_{\mathrm{rev}_i}\),其中 \(\mathrm{rev}_i\) 表示将 \(n\) 位二进制数 \(i\) 的位反转得到的数,则第 \(h\) 层就是对:

\[(\left<i2^{n-h+1}\cdots i2^{n-h+1}+2^{n-h}-1\right>,\left<i2^{n-h+1}+2^{n-h}\cdots (i+1)2^{n-h+1}-1\right>) \]

进行合并。其中 \(0\le i< 2^{h-1}\)。则我们从底向上模拟即可。

实现的时候为了缩小常数我们可以预处理一下单位根。

void FFT(comp* f)
{
	for (int i = 0; i < lim; ++i) if (i < rev[i]) std::swap(f[i], f[rev[i]]);
	for (int h = 2, t = 1, p = lim / h; h <= lim; h <<= 1, t <<= 1, p >>= 1)
	{
		for (int j = 0; j < lim; j += h)
			for (int q = 0, k = j; q < t; ++q, ++k)
			{
				comp u = f[k], v = f[k + t] * wn[p * q];
				f[k] = u + v; f[k + t] = u - v;
			}
	}
}

其中 wn[i] 表示 \(\omega_{\rm lim}^i\)。所以 wn[p * q]\(\omega_{\rm lim}^{\rm q\frac{lim}{h}}=\omega_{\rm h}^{\rm q}\)。这里如果是在模 \(p\) 意义下,\(n\) 次单位根就变成了 \(g^{\frac{p-1}{n}}\),其中 \(g\)\(p\) 的原根。单位根存在,当且仅当 \(n|p-1\)。这也解释了为什么 \(\rm NTT\) 的模数必须满足 \(k2^n+1\) 的形式。

然后我们来看看 \(\rm Bluestein's\ algorithm\)。考虑我们刚刚 \(\rm DFT\) 的过程:

\[\begin{aligned} \hat{f}_i&=\sum_{j=0}^{n-1}f_j\omega_n^{ij}\\ &=\sum_{j=0}^{n-1}f_j\omega_{n}^{\binom{i+j}{2}-\binom{i}{2}-\binom{j}{2}}\\ &=\omega_n^{-\binom{i}{2}}\sum_{j=0}^{n-1}f_j\omega_n^{-\binom{j}{2}}\omega_n^{\binom{i+j}{2}} \end{aligned}\]

其中利用了 \(ij=\binom{i+j}{2}-\binom{i}{2}-\binom{j}{2}\) 这个等式。从而得到了一个差卷积的形式,可以用 \(\rm FFT\) 计算。\(\rm IDFT\) 的过程是类似的。

FFT 实现的优化

利用转置原理对 \(\rm FFT\) 实现进行的常数优化是非常巨大的,且对于减少码量也有帮助。粗体大写表示矩阵,粗体小写表示向量。本部分不会详细系统介绍转置原理。

首先要知道转置定理:

\[\mathbf{C}=\mathbf{AB}\iff \mathbf{C}^{\rm T}=\mathbf{B}^{\mathrm{T}}\mathbf{A}^{\rm T} \]

而我们知道,\(\rm DFT\) 的本质是线性变换,所以有:

\[\mathbf{a}=\mathbf{A}\mathbf{b}\iff \mathbf{a}^{\rm T}=\mathbf{b}^{\rm T}\mathbf{A}^{\rm T} \]

所以如果我们能找到变换矩阵 \(\bf A\) 的转置,就能以另一种方法计算该线性变换。

对于 \(\rm FFT\),它的变换矩阵是范德蒙德矩阵,就像我们看到的,它是关于主对角线对称的,所以它的转置等于自身。也就是说转置完等于没转置。看来从整体入手的优化失效了。

考虑把整个大矩阵分成若干小矩阵,即设:

\[\mathbf{A}=\mathbf{A}_1\mathbf{A}_2\mathbf{A}_3\cdots\mathbf{A}_n \]

则:

\[\mathbf{A}^{\rm T}=\mathbf{A}_n^{\rm T}\mathbf{A}_{n-1}^{\rm T}\mathbf{A}_{n-2}^{\rm T}\cdots\mathbf{A}_1^{\rm T} \]

即我们可以把整个过程分成若干小的线性变换,然后倒着执行每个的转置过程。

对应到 \(\rm FFT\) 上,观察我们的过程:

  • 对序列做 \(\rm bitrev\) 变换,为模拟递归做准备。
  • 模拟递归第 \(n\) 层。
  • 模拟递归第 \(n-1\) 层。
  • \(\vdots\)
  • 模拟递归第 \(1\) 层。

倒过来就是:

  • 模拟递归第 \(1\) 层的转置过程。
  • \(\vdots\)
  • 模拟递归第 \(n\) 层的转置过程。
  • 对序列做 \(\rm bitrev\) 变换的转置过程,这个转置过来等于自身。

如果把它和正常的 \(\rm IDFT\) 过程连起来,就能发现 \(\rm bitrev\) 的过程被消掉了。而事实上,由于内存访问不连续等玄学原因,\(\rm bitrev\) 是比较费时间的。从而在这里优化了常数。现在我们需要解决的问题是,模拟递归的转置过程是什么?

考虑模拟递归对应的矩阵。第 \(i\) 层对应的是一个 \(2^{n-i+1}\times 2^{n-i+1}\) 的矩阵。以 \(n-i+1=3\) 为例:

\[\begin{bmatrix} 1&0&0&0&1&0&0&0\\ 0&1&0&0&0&1&0&0\\ 0&0&1&0&0&0&1&0\\ 0&0&0&1&0&0&0&1\\ \omega_8^{0}&0&0&0&-\omega_8^0&0&0&0\\ 0&\omega_8^1&0&0&0&-\omega_8^1&0&0\\ 0&0&\omega_8^2&0&0&0&-\omega_8^2&0\\ 0&0&0&\omega_8^3&0&0&0&-\omega_8^3 \end{bmatrix}\]

转置!

\[\begin{bmatrix} 1&0&0&0&\omega_8^{0}&0&0&0\\ 0&1&0&0&0&\omega_8^1&0&0\\ 0&0&1&0&0&0&\omega_8^2&0\\ 0&0&0&1&0&0&0&\omega_8^3\\ 1&0&0&0&-\omega_8^0&0&0&0\\ 0&1&0&0&0&-\omega_8^1&0&0\\ 0&0&1&0&0&0&-\omega_8^2&0\\ 0&0&0&1&0&0&0&-\omega_8^3 \end{bmatrix}\]

这样就得到了转置之后的式子。

\[\hat{f}_i=\hat{f}_{0,i}+\hat{f}_{1,i},\hat{f}_{i+n/2}=\omega_n^i(\hat{f}_{0,i}-\hat{f}_{1,i})(i<n/2) \]

从而可以写出代码。

inline void NTT(int* f)
{
	for (int h = lim, t = lim >> 1, p = 1; h >= 2; h >>= 1, t >>= 1, p <<= 1)
		for (int j = 0; j < lim; j += h)
			for (int k = j, q = 0; q < t; ++q, ++k)
			{
				int u = f[k], v = f[k + t];
				f[k] = (u + v) % mod; f[k + t] = (ll)w[0][p * q] * (u + mod - v) % mod;
			}
}
inline void INTT(int* f)
{
	for (int h = 2, t = 1, p = (lim >> 1); h <= lim; h <<= 1, t <<= 1, p >>= 1)
		for (int j = 0; j < lim; j += h)
			for (int k = j, q = 0; q < t; ++q, ++k)
			{
				int u = f[k], v = (ll)w[1][p * q] * f[k + t] % mod;
				f[k] = (u + v) % mod; f[k + t] = (u + mod - v) % mod;
			}
	for (int i = 0, inv = ksm(lim, mod - 2); i < lim; ++i) f[i] = (ll)f[i] * inv % mod;
}

其中 w[0/1][i] 分别表示 \(\omega_{\rm lim}^{i}\)\(\omega_{\rm lim}^{-i}\)。真的很快还好写

写完了。我大概是想不开了才会在 \(\rm NOI\) 前总结这东西。

posted @ 2022-08-10 16:13  zhiyangfan  阅读(236)  评论(0编辑  收藏  举报