类 FFT 状物笔记(持续更新)

· · 算法·理论

?!积积!?

记录了 Hootime 学过的所有类 FFT 状物。如果 Hootime 学了更多的类 FFT 状物,它会来修改这篇文章。另外你也可以在 cnblog 上阅读这篇文章。

推导的时候有一定的跳步。

共通性 / Commons

都是这么一个问题:

给定两个序列 A, B,你需要求出一个序列 C,其中 C_i = \sum_{j \oplus k = i} A_j B_k。你需要一个 \mathrm O(n \log n) 的做法。

都是这么一个做法。

构造一个序列 f(A),其中将 A 变换到 f(A) 或者将 f(A) 变换到 An \log n 的,且 f(C)_{i} = f(A)_{i} f(B)_{i}

只是不同的问题有不同的做法。

快速傅里叶变换 / FFT

需要解决的问题是:C_i = \sum_{j + k = i} A_j B_k

我们将序列 A 设为一个 n 次函数,其中 g_A(x) = \sum_{i=0}^n A_ix^i。发现 g_C(x) = g_A(x) g_B(x),于是我们取 nx 就满足了上一节提到的 f(A) 的第二条性质。

问题在于这个转换是 \mathrm O(n^2) 的。

离散傅里叶变换 / DFT

指将 A 变换为 f(A) 的过程。

以下定义 \omega_n^k 为第 kn 次单位根。

以下设序列长度 n = 2^m。如果 n 不是 2 的幂次,可以通过补 0 让它是。序列是 0-index 的。

我们观察以下这个式子。以下将 g_A(x) 称为 g(x)

\begin{aligned} g(x) &= A_0x^0 + A_1x^1 + A_2x^2 + A_3x^3 + \cdots + A_{n-1}x^{n-1} + A_nx^n \\ &= (A_0x^0 + A_2x^2 + \cdots A_{n-1}x^{n-1}) + (A_1x^1 + A_3x^3 + \cdots + A_nx^n) \\ &= (A_0x^0 + A_2x^2 + \cdots A_{n-1}x^{n-1}) + x(A_1x^0 + A_3x^2 + \cdots + A_nx^{n-1}) \end{aligned}

然后我们定义:

\begin{aligned} g_0(x) &= A_0x^0 + A_2x^1 + \cdots + A_{n-1}x^{n/2-1} \\ g_1(x) &= A_1x^0 + A_3x^1 + \cdots + A_nx^{n/2-1} \end{aligned}

于是发现:

g(x) = g_0(x^2)+xg_1(x^2)

于是我们考虑将 \omega_n^0, \omega_n^1, \cdots, \omega_n^{n-1} 作为我们选择的 nx

我们将 \omega_n^k (k \le \frac n 2) 代进去。于是:

\begin{aligned} g(\omega_n^k) &= g_0(\omega_n^{2k})+\omega_n^kg_1(\omega_n^{2k}) \\ &= g_0(\omega_{n/2}^k)+\omega_n^kg_1(\omega_{n/2}^k) \end{aligned}

同理,我们将 \omega_n^{k + n/2} (k \le \frac n 2) 代进去。于是:

\begin{aligned} g(\omega_n^{k+n/2}) &= g_0(\omega_n^{2(k+n/2)})+\omega_n^{k+n/2}g_1(\omega_n^{2(k+n/2)}) \\ &= g_0(\omega_{n/2}^k)-\omega_n^kg_1(\omega_{n/2}^k) \end{aligned}

于是我们发现 g_0, g_1g 本质相同,我们递归一下即可。

然后递归常数太大了,我们要删掉递归。你可以试试不删,你不一定能过 luogu 的板子

我们思考一下有没有迭代的做法。我们观察一下递归的顺序,发现:

于是惊奇地发现,在递归时被分在一组的位置在二进制翻转后是相邻的(翻转指顺序翻转而非 0/1 翻转为 1/0)!于是我们就可以直接迭代了。

于是我们完成了 A \rightarrow f(A) 的转换。

inline void dft(cdouble a[]){
    for(int i = 0; i < n; ++i)
        if(i < rev[i]) swap(a[i], a[rev[i]]);
    for(int l = 2; l <= n; l <<= 1){
        cdouble wn = {cos(2*pi/l), sin(2*pi/l)};
        for(int i = 0; i < n; i += l){
            cdouble w = {1, 0};
            for(int j = i; j < i+l/2; ++j){
                cdouble res1 = a[j], res2 = w*a[j+l/2];
                a[j] = res1+res2, a[j+l/2] = res1-res2;
                w = w*wn;
            }
        }
    }
    return;
}

其中 rev[i]i 翻转后的结果,cdoublecomplex<double>

离散傅里叶逆变换 / IDFT

接下来我们解决另一部分,即将 f(A) 变换为 A我不知道第一个独立构造出来的人是怎么做到的但是真的强啊

我们接下来记 f(A)_ih_i。于是我们直接将 h 视作一个多项式,构造 g(x) = \sum_{i=0}^{n-1} h_ix^i,并且代入 \omega_n^0, \omega_n^{-1}, \cdots, \omega_n^{-n+1}

下面我们来推导 g(x)

\begin{aligned} g(\omega_n^k) &= \sum_{i=0}^{n-1} h_i\omega_n^{-ik} \\ &= \sum_{i=0}^{n-1} (\sum_{j=0}^{n-1} a_j \omega_n^{ij}) \omega_n^{-ik} \\ &= \sum_{i=0}^{n-1} \sum_{j=0}^{n-1} a_j \omega_n^{i(j-k)} \\ &= \sum_{j=0}^{n-1} a_j (\sum_{i=0}^{n-1} (\omega_n^{j-k})^i) \end{aligned}

然后我们令 S = \sum_{i=0}^{n-1} (\omega_n^{k})^i

k \neq 0 时,根据等比数列的推导过程:

\begin{aligned} S &= \sum_{i=0}^{n-1} (\omega_n^k)^i \\ \omega_n^k S &= \sum_{i=1}^{n} (\omega_n^k)^i \\ (\omega_n^k-1) S &= \omega_n^k - \omega_n^0 \\ &= 0 \\ S &= 0 (\omega_n^k \neq 0) \end{aligned}

k = 0 时,显然 S = n。也就是说:

S = \begin{cases} n & \text{if }k = 0 \\ 0 & \text{otherwise.} \end{cases}

也就是说:

\begin{aligned} g(\omega_n^k) &= \sum_{j=0}^{n-1} a_j (\sum_{i=0}^{n-1} (\omega_n^{j-k})^i) \\ &= na_k \end{aligned}

也就是说 a_k = g(\omega_n^{-k}) \cdot n^{-1}

于是我们完成了 f(A)A 的转换。我们发现两个转换可以合并成一个。

inline void dft(cdouble a[], int inv){
    for(int i = 0; i < n; ++i)
        if(i < rev[i]) swap(a[i], a[rev[i]]);
    for(int l = 2; l <= n; l <<= 1){
        cdouble wn = {cos(2*pi/l), inv*sin(2*pi/l)};
        for(int i = 0; i < n; i += l){
            cdouble w = {1, 0};
            for(int j = i; j < i+l/2; ++j){
                cdouble res1 = a[j], res2 = w*a[j+l/2];
                a[j] = res1+res2, a[j+l/2] = res1-res2;
                w = w*wn;
            }
        }
    }
    if(inv == -1)
        for(int i = 0; i < n; ++i) a[i] /= n;
    return;
}

快速数论变换 / NTT

需要解决的问题是:C_i = (\sum_{j + k = i} A_j B_k) \bmod p,其中 p 有原根 c

我们当然可以使用 FFT 然后取模,但是你精度炸了。

我们将单位根的 k 次方替换成原根的 k 次方,将除以 n 替换成乘以 n 的逆元,于是你就将 FFT 替换成了 NTT。所以其实 FFT 和 NTT 本质相同。

constexpr llong p = 998244353;
inline void ntt(llong a[], int inv){
    for(int i = 0; i < n; ++i)
        if(i < rev[i]) swap(a[i], a[rev[i]]);
    for(int l = 2; l <= n; l <<= 1){
        llong wn = qpow(3, (p-1)/l);
        if(inv == -1) wn = qpow(wn, p-2);
        for(int i = 0; i < n; i += l){
            llong w = 1;
            for(int j = i; j < i+(l>>1); ++j){
                llong r1 = a[j], r2 = w*a[j+(l>>1)]%p;
                a[j] = (r1+r2)%p, a[j+(l>>1)] = (r1-r2+p)%p;
                w = w*wn%p;
            }
        }
    }
    if(inv == -1){
        int invn = qpow(n, p-2);
        for(int i = 0; i < n; ++i) a[i] = a[i]*invn%p;
    }
    return;
}

注意因为我是唐人,这份代码是 \mathrm O(n \log^2 n) 的。如果需要改成 n \log n 可以把循环里面的 qpow 放在外面,但是我懒。

我们考虑 p 没有性质的情况,这是我们多选几个模数做 NTT 然后 CRT 合并一下子最后取模即可。发现答案值域是 nV^2 的(VA 的值域),所以一般选三个。

快速沃尔什变换 / FWT

需要解决的问题是:C_i = \sum_{j \otimes k = i} A_j B_k,其中 \otimes\mathrm{and}, \mathrm{or}, \mathrm{xor} 中的一种。

OR 卷积

我们还是考虑构造 f(A)

直接给出构造:f(A)_i = \sum_{j \operatorname{or} i = i} A_j。下证这样的 f(A) 满足性质。

首先显然地,若 j \operatorname{or} i = i, k \operatorname{or} i = i 那么 (j \operatorname{or} k) \operatorname{or} i = i

那么有:

\begin{aligned} f(A)_if(B)_i &= (\sum_{j \operatorname{or} i = i} A_j)(\sum_{k \operatorname{or} i = i} B_k) \\ &= \sum_{j \operatorname{or} i = i} \sum_{k \operatorname{or} i = i} A_j B_k \\ &= \sum_{(j \operatorname{or} k) \operatorname{or} i = i} A_j B_k \\ &= f(C)_i \end{aligned}

于是构造的正确性得证。

于是我们可以使用高维前缀和来实现 A \rightarrow f(A)、用高维差分来实现 f(A) \rightarrow A。但是这是一个类 FFT 状物,所以我们有一个与 FFT 类似且等价于高维前缀和的表示方法:(记 f(A)_ig(i)

\begin{aligned} g(i) &= g_0(i) \\ g(i+\frac n 2) &= g_1(i)+g_0(i) \end{aligned}

同样地,我们可以表示逆变换:

\begin{aligned} g(i) &= g_0(i) \\ g(i+\frac n 2) &= g_1(i)-g_0(i) \end{aligned}
constexpr llong p = 998244353;
inline void _or(llong* a, int inv){
    for(int len = 2; len <= n; len <<= 1)
        for(int i = 0; i < n; i += len)
            for(int j = 0; j < (len>>1); ++j)
                (a[i+j+(len>>1)] += inv*a[i+j]) %= p;
    return;
}

其中,inv = 1 代表正变换,inv = p-1 代表逆变换。w

AND 卷积

等价于 OR 卷积。

inline void _and(llong* a, int inv){
    for(int len = 2; len <= n; len <<= 1)
        for(int i = 0; i < n; i += len)
            for(int j = 0; j < (len>>1); ++j)
                (a[i+j] += inv*a[i+j+(len>>1)]) %= p;
    return;
}

XOR 卷积

我们还是考虑构造 f(A)

还是直接给出构造:定义 a \oplus b = \mathrm{popcount}(a \operatorname{and} b) \bmod 2,那么 f(A)_i = \sum_{j \oplus i = 0}A_j - \sum_{j \oplus i = 1}A_j

还是给出证明:

\begin{aligned} f(A)_if(B)_i &= (\sum_{j \oplus i = 0}A_j - \sum_{j \oplus i = 1}A_j)(\sum_{k \oplus i = 0}B_k - \sum_{k \oplus i = 1}B_k) \\ &= (\sum_{j \oplus i = 0}A_j)(\sum_{k \oplus i = 0}B_k) - (\sum_{j \oplus i = 1}A_j)(\sum_{k \oplus i = 0}B_k) - (\sum_{j \oplus i = 0}A_j)(\sum_{k \oplus i = 1}B_k) + (\sum_{j \oplus i = 1}A_j)(\sum_{k \oplus i = 1}B_k) \\ &= \sum_{(j \operatorname{xor} k) \oplus i = 0} A_jB_k - \sum_{(j \operatorname{xor} k) \oplus i = 1} A_jB_k \end{aligned}

于是得证。

但是怎么变换呢?

我们还是考虑类似 FFT 地,将 f(A) 切成 [0, \frac n 2), [\frac n 2, n) 两段。发现 [0, \frac n 2), [\frac n 2, n) 的最高位分别是 01。又发现若 i \operatorname{and} j = 1(i, j \in \{ 0, 1 \}),当且仅当 i = j = 1,那么发现 g_1g 的贡献要乘以 -1,其它不变。

那么:

\begin{aligned} g(i) &= g_0(i)+g_1(i) \\ g(i+\frac n 2) &= g_0(i) - g_1(i) \end{aligned}

逆变换就是:

\begin{aligned} g(i) &= \frac {g_0(i)+g_1(i)} 2 \\ g(i+\frac n 2) &= \frac {g_0(i) - g_1(i)} 2 \end{aligned}

也就是说:

inline void _xor(llong* a, int inv){
    for(int len = 2; len <= n; len <<= 1){
        for(int i = 0; i < n; i += len){
            for(int j = 0; j < (len>>1); ++j){
                llong a0 = a[i+j], a1 = a[i+j+(len>>1)];
                a[i+j]          = (a0+a1+p)*inv%p;
                a[i+j+(len>>1)] = (a0-a1+p)*inv%p;
            }
        }
    }
    return;
}

其中正变换时 inv = 1,逆变换时 inv = (p+1)/2

注意对于我的写法,写 FWT 模板需要把数组开到 1<<17|5 而非 1<<17,或者关掉 O2。别问,问就是 SIMD 发力了。