📝 상세 정리

• 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$ 을 사용할 수 있다.
  • 그러면 이제 $\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;
	}
}

❔질문 사항

🔗 참고 자료

• https://rkm0959.tistory.com/187
• https://algoshitpo.github.io/2020/05/20/fft-ntt/