number theoretic transform

目的

Z/(pZ)上の係数を持つ$n$次多項式f(x),g(x)の積f(x)g(x)を$O(n \log n)$で求めます。通常の高速フーリエ変換に比べて、計算誤差がない・計算時間が短いという利点があります。

計算量

$O(n\log n)$

使い方

long[] mul(long[] a,long[] b)
f(x)=a[i]x^i,g(x)=b[i]x^iとしたときのf(x)g(x)のi次の係数がi番目に格納された配列をreturnします。
long MODULO
目的の項の$p$に相当します。
long root
Z/(pZ)の原子根です。

ソースコード

long[] mul(long[] a, long[] b) {
	int n = Integer.highestOneBit(a.length + b.length) << 1;
	a = Arrays.copyOf(a, n);
	b = Arrays.copyOf(b, n);
	a = fft(a, false);
	b = fft(b, false);
	long ninv = pow_mod(n, MODULO - 2);
	for (int i = 0; i < n; ++i) {
		a[i] = a[i] * b[i] % MODULO;
	}
	a = fft(a, true);
	for (int i = 0; i < n; ++i)
		a[i] = a[i] * ninv % MODULO;
	return a;
}
long[] fft(long[] a, boolean inv) {
	int n = a.length;
	int c = 0;
	for (int i = 1; i < n; ++i) {
		for (int j = n >> 1; j > (c ^= j); j >>= 1)
			;
		if (c > i) {
			long d = a[c];
			a[c] = a[i];
			a[i] = d;
		}
	}
	for (int i = 1; i < n; i <<= 1) {
	long w = pow_mod(root, (MODULO - 1) / (2 * i));
	if (inv)
		w = pow_mod(w, MODULO - 2);
	for (int j = 0; j < n; j += 2 * i) {
		long wn = 1;
		for (int k = 0; k < i; ++k) {
			long u = a[k + j];
				long v = a[k + j + i] * wn % MODULO;
				a[k + j] = (u + v);
				if (a[k + j] >= MODULO)
					a[k + j] -= MODULO;
				if (a[k + j] < 0)
					a[k + j] += MODULO;
				a[k + j + i] = u - v;
				if (a[k + j + i] >= MODULO)
					a[k + j + i] -= MODULO;
				if (a[k + j + i] < 0)
					a[k + j + i] += MODULO;
				wn = wn * w % MODULO;
			}
		}
	}
	return a;
}

long pow_mod(long a, long n) {
	long ret = 1;
	for (; n > 0; n >>= 1, a = a * a % MODULO) {
		if (n % 2 == 1) {
			ret = ret * a % MODULO;
		}
	}
	return ret;
}

long inv(long a) {
	return pow_mod(a, MODULO - 2);
}