Fourier 変換ちゃん

[rsk0315]

• https://satanic0258.github.io/snippets/math/FFT.html
• https://judge.yosupo.jp/submission/63348

さすがに速い FFT たちの中に Karatsuba 法だけで立ち向かうのはつらいので

コード

```rs use proconio::{fastout, input}; use nekolib::math::GcdRecip; const MOD: u64 = 998244353; #[fastout] fn main() { input! { n: usize, m: usize, a: [u64; n], b: [u64; m], } let c = convolution_mod(a, b, MOD); for &ci in &c { println!("{}", ci); } } fn convolution_mod(mut a: Vec, mut b: Vec, p: u64) -> Vec { if a.is_empty() && b.is_empty() { return vec![]; } let m = a.len() + b.len() - 1; let n = m.next_power_of_two(); a.resize(n, 0); b.resize(n, 0); let g = { let mut g = omega(p); for _ in n.trailing_zeros()..(p - 1).trailing_zeros() { g *= g; g %= p; } g }; let omega: Vec<_> = std::iter::successors(Some(1), |&x| Some(x * g % p)).take(n).collect(); let h = g.gcd_recip(MOD).1; let omega_inv: Vec<_> = std::iter::successors(Some(1), |&x| Some(x * h % p)).take(n).collect(); let fa = ntt(a, &omega_inv, p); let fb = ntt(b, &omega_inv, p); let fab: Vec<_> = fa.into_iter().zip(fb).map(|(x, y)| x * y % MOD).collect(); let mut ab = ntt(fab, &omega, p); ab.resize(m, 0); let n_inv = (n as u64).gcd_recip(MOD).1; for e in ab.iter_mut() { *e = *e * n_inv % MOD; } ab } fn ntt(mut a: Vec, omega: &[u64], p: u64) -> Vec { let n = a.len(); let mask1 = n - 1; let mut tmp = vec![0; n]; for i in 0..n.trailing_zeros() { let mask2 = mask1 >> (i + 1); for j in 0..n { let lo = j & mask2; let hi = j ^ lo; let sh = (hi << 1) & mask1; tmp[j] = (a[sh | lo] + omega[hi] * a[sh | (mask2 + 1) | lo]) % p; } std::mem::swap(&mut a, &mut tmp); } a } fn omega(p: u64) -> u64 { match p { 998244353 => 15311432, // 3 ^ 119 mod p _ => todo!(), } } ```

それはそれとして、いい感じの mod int とか、実行時 mod はどうしようとか、そういうのを考え出すと大変になる

[rsk0315]

実行時 mod int で <0> は 10000、<1> は 1234、みたいなのをやろうとすると、どうしても lazy_static! 的なのが欲しくなるはず

[rsk0315]

そもそも FFT の interface 自体困る とりあえず上記のは、f(x) = a[0] + a[1] x + a[2] x^2 + ... + a[n - 1] x^(n - 1) と mod p での 1 の原始 n 乗根 omega に対して

• (0..n).map(|i| a[i]).collect()
• (0..n).map(|i| omega.gcd_recip(p).1.mod_pow(i, p)).collect()
• p

を渡すことにしてみた。返すのは sum(a[j] omega ^ (i * j)) だけど、渡すのは omega の逆数の累乗なのでややこしいかも？ 逆向きにアクセスすることにすれば omega の累乗でもいいけど、長さが n + 1 になるから微妙そう

[rsk0315]

omega を渡すようにしたのは、どうせ畳み込みするときに同じものを 2 回使う必要があるからなんだけど、そもそも別で struct を作るとかの方がよかったりする？ それもなんか微妙な気がするんだよねえ