📝 상세 정리#
- FFT를 잘 모른다면, 먼저 읽고 오자.
- 다항식의 곱셈 결과를 특정 소수 $p$로 나눈 나머지를 구해야하거나, 값이 커지거나 해서 실수 오차가 발생하는 경우가 있다.
- 실수를 안쓰고 FFT를 할 수는 없을까?
- 우리가 FFT에서 핵심으로 사용하던 $\omega = \exp\left(-i\dfrac{2\pi}{N}\right)$를 조금 더 생각해보자.
- FFT에서 사용되는 핵심적인 성질은 $\omega^N = 1$ 이고, $1, \omega, \omega^2, \cdots, \omega^{N-1}$이 모두 다르다는 점이다.
- 그러면 모듈러 합동식의 관점에서, 어떤 수 $k$와 어떤 소수 $p$가 있어서 $k^N \equiv 1 \pmod p$를 만족시키고, $1, g, g^2, \cdots, g^{N-1} \pmod p$이 모두 다르다면, FFT와 같은 느낌으로 사용할 수 있지 않을까?
- 이러한 수 $g$를 원시근이라고 한다.
- 그러면 $N$을 또 잘 잡아야할 것 같은데, 페르마의 소정리를 이용할 수 있을 것 같다.
- $a^{p-1} \equiv 1 \pmod p$ 이므로, $N = p-1$이면 좋을 것 같다.
- FFT를 반으로 쪼개던 방법을 잘 이용하기 위해선, $N$이 $2$로 충분히 많이 나누어 떨어지면 좋겠다.
- $N = a \times 2^b$ 꼴이면 좋겠고, 이에 따라 $p = a \times 2^b + 1$ 이면 좋겠다.
- $N$을 $2$로 나누어 떨어지게 할 수 있는 횟수에 따라 FFT로 계산 가능한 배열의 크기가 달라진다!
- $n \leq 2^b$ 의 크기의 다항식까지만 가능하다.
- 이로 어울리는 대표적인 소수로는 $p = 998244353 = 119 \times 2^{23} + 1$ 이 있겠다.
- 이때 원시근 $g = 3$ 을 사용할 수 있다.
- $N = a \times 2^b$ 꼴이면 좋겠고, 이에 따라 $p = a \times 2^b + 1$ 이면 좋겠다.
- 그러면 이제 $\omega = \exp\left(-i\dfrac{2\pi}{N}\right)$ 대신 $\omega = g^{(p-1)/n}$ 을 이용해서 FFT를 수행할 수 있게 되었다.
- 이를 구현하면 다음과 같다.
void ntt(vector<mint> &a, bool inv = false){
int N = a.size();
assert(N > 0 && (N & (N-1)) == 0); // N은 2의 거듭제곱이어야 함
// 비트 뒤집기
for(int i = 1, j = 0; i < N; i++){
int bit = N >> 1;
for(; j & bit; bit >>= 1) j ^= bit;
j ^= bit;
if(i < j) swap(a[i], a[j]);
}
// 단위근 미리 계산
mint ang = mint(3).pow((MOD-1)/N);
if(inv) ang = ang.inv();
vector<mint> w(N/2); // 단위근
w[0] = 1;
rep(i, 1, N/2) w[i] = w[i-1] * ang;
// Cooley-Tukey FFT
for(int len = 2; len <= N; len <<= 1){
int step = N / len;
for(int i = 0; i < N; i += len){
rep(j, 0, len>>1){
mint even = a[i+j], odd = a[i+j+(len>>1)] * w[j*step];
a[i+j] = even + odd;
a[i+j+(len>>1)] = even - odd;
}
}
}
// 역 FFT인 경우 결과를 N으로 나눔
if(inv){
mint invN = mint(N).inv();
rep(i, 0, N) a[i] *= invN;
}
}