Number Theoretic Fast Fourier Transform

2^n 크기의 배열의 FFT를 구할 때, 우리는 다음과 같은 복소수 w를 이용한다.

w = cos(2π/k) + i sin(2π/k)

여기서 우리가 주목할 점은, k는 2의 거듭제곱이며 w^k = 1 이라는 사실이다.

그러므로 우리는 다음과 같은 꼴을 가지는 소수 p를 생각해 볼 수 있다.

p = a * 2^b + 1

이런 p를 잡으면 p의 원시근 x에 대하여 (x^a)^(2^b) mod p = 1 이고, 따라서 w 대신 (x^a)^(2^b/k)를 사용하여 n<=b를 만족하는 모든 2^n 크기의 배열에 대해 법 p로 FFT를 행할 수 있다.

아래는 위를 만족하는 충분히 큰 소수들 목록이다.

p	a	b	원시근	덧셈	곱셈
3221225473	3	30	5	64bit signed	64bit unsigned
2281701377	17	27	3	64bit signed	64bit signed
2013265921	15	27	31	32bit unsigned	64bit signed
469762049	7	26	3	32bit signed	64bit signed

각 수마다 덧셈과 곱셈을 행할 때 사용해야 하는 자료형이 다르므로 상황에 맞게 사용하면 좋다.

#include <cstdio>

const int A = 7, B = 26, P = A << B | 1, R = 3;
const int SZ = 20, N = 1 << SZ;

int Pow(int x, int y) {
	int r = 1;
	while (y) {
		if (y & 1) r = (long long)r * x % P;
		x = (long long)x * x % P;
		y >>= 1;
	}
	return r;
}

void FFT(int *a, bool f) {
	int i, j, k, x, y, z;
	j = 0;
	for (i = 1; i < N; i++) {
		for (k = N >> 1; j >= k; k >>= 1) j -= k;
		j += k;
		if (i < j) {
			k = a[i];
			a[i] = a[j];
			a[j] = k;
		}
	}
	for (i = 1; i < N; i <<= 1) {
		x = Pow(f ? Pow(R, P - 2) : R, P / i >> 1);
		for (j = 0; j < N; j += i << 1) {
			y = 1;
			for (k = 0; k < i; k++) {
				z = (long long)a[i | j | k] * y % P;
				a[i | j | k] = a[j | k] - z;
				if (a[i | j | k] < 0) a[i | j | k] += P;
				a[j | k] += z;
				if (a[j | k] >= P) a[j | k] -= P;
				y = (long long)y * x % P;
			}
		}
	}
	if (f) {
		j = Pow(N, P - 2);
		for (i = 0; i < N; i++) a[i] = (long long)a[i] * j % P;
	}
}

int X[N];

int main() {
	int i, n;
	scanf("%d", &n);
	for (i = 0; i <= n; i++) scanf("%d", &X[i]);
	FFT(X, false);
	for (i = 0; i < N; i++) X[i] = (long long)X[i] * X[i] % P;
	FFT(X, true);
	for (i = 0; i <= n + n; i++) printf("%d ", X[i]);
}