算法 - FWT & FFT & NTT

Posted on Edited on In Disqus:

反正是用来卷的。

FWT

作用

用于求 \(C[k]=\sum_{g(i,j)=k} A[i] \cdot B[j]\)，其中 \(g(i,j)\) 是一个按位做一个二元运算的函数。

开始

记 \(C=f(A,B)\)，\(A=A_0::A_1\)，其中 \(::\) 表示连接， \(A_0,A_1\) 分别表示 A 的前一半和后一半（首位分别为 0, 1）。

拿 g 为 XOR 为例： \[ C_0 = f(A_0, B_0) + f(A_1, B_1) \\ C_1 = f(A_0, B_1) + f(A_1, B_0) \\ C_0 + C_1 = f(A_0, B_0) + f(A_1, B_1)+f(A_0, B_1) + f(A_1, B_0)=f(A_0+B_0,A_1+B_1) \\ C_0 - C_1 = f(A_0, B_0) + f(A_1, B_1)-f(A_0, B_1) - f(A_1, B_0)=f(A_0-B_0,A_1-B_1) \\ (C_0 + C_1) :: (C_0 - C_1) = f(A_0+B_0,A_1+B_1)::f(A_0-B_0,A_1-B_1) \] 拿 g 为 OR 为例： \[ C_0 = f(A_0, B_0) \\ C_1 = f(A_0, B_1) + f(A_1, B_0) + f(A_1, B_1) \\ C_0 = f(A_0, B_0) \\ C_0 + C_1 = f(A_0, B_0) + f(A_0, B_1) + f(A_1, B_0) + f(A_1, B_1)=f(A_0+B_0,A_1+B_1) \\ C_0::(C_0 + C_1) = f(A_0, B_0)::f(A_0+B_0,A_1+B_1) \] 我们递归地定义 \(fwt\) 函数，边界为 \(fwt(a)=a\)。

对于 XOR，\(fwt(A)=fwt(A_0::A_1)=fwt(A_0+A_1)::fwt(A_0-A_1)\)。

对于 OR，\(fwt(A)=fwt(A_0::A_1)=fwt(A_0)::fwt(A_0+A_1)\)。

利用之前推出的等式，可以归纳证明 \(fwt(C)=fwt(A) \circ fwt(B)\)，其中 \(\circ\) 为对应位相乘。

那么 \(C=fwt^{-1}(fwt(A)\circ fwt(B))\)，至此 \(fwt\) 完成了。

以 XOR 为例简单证明，设 \(A,B,C\) 中的元素个数为 \(n\)。

\(n=1\) 时显然成立。

若 \(n=k\) 时成立，即对于 \(C=f(A,B)\) 有 \(fwt(C)=fwt(A) \circ fwt(B)\)。

\(n=2k\) 时 \[ \begin{eqnarray} &&fwt(A)\circ fwt(B) \\ &=&fwt(A_0::A_1)\circ fwt(B_0::B_1) \\ &=&(fwt(A_0+A_1)::fwt(A_0-A_1))\circ(fwt(B_0+B_1)::fwt(B_0-B_1)) \\ &=& (fwt(A_0+A_1) \circ fwt(B_0+B_1)) :: (fwt(A_0-A_1) \circ fwt(B_0-B_1)) \\ &=& fwt(f(A_0+A_1,B_0+B_1))::fwt(f(A_0-A_1,B_0-B_1)) \\ &=& fwt(C_0+C_1)::fwt(C_0-C_1) \\ &=& fwt(C_0::C_1) \\ &=& fwt(C) \end{eqnarray} \]

实现

  • 如果这个运算 g 不满足交换律，你会发现类似 \(fwt(C)=fwt(A) \circ fwt'(B)\) 的式子，一种方法就是多写一个函数，或者把 A 或 B 取反，那么 g 就满足交换律了。
  • FWT 可以简单地写成非递归的形式。
  • 对于 XOR，有 \(fwt^{-1}(A)=fwt(A)/n\)，因为每递归一层就会差 2 倍。
1
2
3
4
5
6
7
8
9
10
11
12
13
template<typename T>
void fwt(LL a[], int n, T f) {
    for (int d = 1; d < n; d *= 2)
        for (int i = 0, t = d * 2; i < n; i += t)
            FOR (j, 0, d)
                 f(a[i + j], a[i + j + d]);
}

void XOR (LL& a, LL& b) {
    LL x = a, y = b;
    a = (x + y) % MOD;
    b = (x - y + MOD) % MOD;
}

FFT

作用

已知 \(A,B \in \mathbb R[X]\)，求 \(C=A\cdot B\)（多项式乘法）。

设 \(\deg(P)=n\)，选取 \(n+1\) 个不同的 \(x_i\)，得到该多项式的一个点值表示 \(\{(x_i,P(x_i))\}\)，而一个 \(n+1\) 个点值表示能唯一确定一个不超过 \(n\) 次的多项式。

如果能够获得 \(A,B\) 的点值表示法（对于同样的 \(x_i\)），那么可以在线性时间内求出 \(C\) 的点值表示法。

而 FFT 正是用于加速这个过程。

开始

DFT

为方便计算，设 \(\deg(A)=2^r - 1=N-1, A(x)=\displaystyle \sum_{i=0}^{N-1} a_i x^i\)。

选取的 \(N\) 个 \(x_k\) 为 \(N\) 次单位根，即 \(\omega_n^k=e^{\frac{2\pi ki}{n}}\)。

现在要求 \(F(k)=A(x_k)=A(\omega_n^k)=\sum_n a_i \omega_n^{ki}\)，求奇偶项分别求和，设 \(n=2t\)，得到 \(F(k)=\sum_t a_{2t}\omega_{t}^{kt}+\omega_n^k \sum_ta_{2t+1}\omega_{t}^{kt}=F_{even}(k)+\omega_n^k F_{odd}(k)\)。

对于奇数项和偶数项，分别求出 \(F_{even}\) 和 \(F_{odd}\)，但是由于项数只有一半，所以直接获得的只有前 \(t\) 个点。由于 \(t\) 是 \(w_t^k\) 的周期且有 \(\omega_n^k+\omega_n^{k+t}=0\)，所以有 \(F(k+t)=F_{even}(k)-\omega_n^k F_{odd}(k)\)。

至此 DFT 完成了。

IDFT

但是为了求出 \(C\)，还需要一个 DFT 的逆变换 IDFT。

考虑 DFT 的过程，设矩阵 \(M_{ij}=\omega_n^{ij}\)，那么 \(DFT(\vec{a})=M \cdot \vec{a}\)

现在需要求 \(IDFT(DFT(\vec{a}))=\vec{a} = M^{-1} \cdot DFT(\vec{a})\)。

设 \(M^{-1}=M'\)，那么有 \(M'_{ij}=\frac{1}{n}\omega_n^{-ij}\)。

所以只需要对 \(DFT\) 算法稍作改动就得到了 \(IDFT\)。

至此 FFT 完成了。

实现

  • DFT 由于要奇偶分别递归，要写成非递归形式可能有些困难。目标是把初始向量重新排列之后变成左右递归而不是奇偶递归，以方便自底向上（其实就是把递归的依据从下标的末尾换成首位），这部分怎么写炫酷就怎么写。

NTT

与 FFT 唯一的不同是多项式是模意义下的。（\(\mathbb{Z}_m[X]\)）

此时需要单位根的替代品，原根，但同时对模数提出了一些要求。

但实现方法和 FFT 是一致的。

此处暂略。

后记

鉴于网上的一些资料不够直观/简洁，于是按自己的思路另外写了一份。