fft.convolution

数列の畳み込み

2つの数列 \(A_i, B_i\) について, \[C_k = \sum_{i=0}^k A_i B_{k-i}\] なる数列 \(C_k\) を計算する.

例題

fft.convolution.rs

struct FFT;

impl FFT {

    fn _fft(f: &Vec<Complex>, dir: K) -> Vec<Complex> {
        let n = f.len();
        if n < 2 { return f.clone() }
        let f0: Vec<Complex> = (0..n/2).map(|i| f[i*2].clone()).collect();
        let f1: Vec<Complex> = (0..n/2).map(|i| f[i*2+1].clone()).collect();
        let g0 = FFT::_fft(&f0, dir);
        let g1 = FFT::_fft(&f1, dir);
        let pi = (-1.0 as K).acos();
        let theta = 2.0 * pi / (n as K);
        let z = Complex(theta.cos(), theta.sin() * dir);
        let mut ac = Complex::unit();
        let mut g = vec![];
        for i in 0..n {
            g.push(&g0[i % (n/2)] + &(&ac * &g1[i % (n/2)]));
            ac = &z * &ac;
        }
        g
    }
    fn fft(f: &Vec<Complex>) -> Vec<Complex> { FFT::_fft(f, 1.0) }
    fn defft(f: &Vec<Complex>) -> Vec<Complex> { FFT::_fft(f, -1.0) }

    fn _convolution(x: &Vec<Complex>, y: &Vec<Complex>) -> Vec<Complex> {
        let m = x.len();
        let xh = FFT::fft(&x);
        let yh = FFT::fft(&y);
        let xyh = (0..m).map(|i| &(&xh[i] * &yh[i]) * (1.0 / m as K)).collect();
        FFT::defft(&xyh)
    }

    fn convolution(f: &Vec<K>, g: &Vec<K>) -> Vec<K> {
        let n = max(f.len(), g.len());
        let mut m = 1; while m < n { m <<= 1 }; m <<= 1; // length should be 2**pow
        let x: Vec<Complex> =
            f.iter().map(|&k| Complex(k, 0.0)).chain(
                (0..m - f.len()).map(|_| Complex(0.0, 0.0))).collect();
        let y: Vec<Complex> =
            g.iter().map(|&k| Complex(k, 0.0)).chain(
                (0..m - g.len()).map(|_| Complex(0.0, 0.0))).collect();

        let z = FFT::_convolution(&x, &y);
        z.iter().map(|c| c.0).collect()
    }
}

fft.convolution.cc

using Real = long double;
const Real PI = acos(-1);
using C = complex<Real>;

vector<C>
fft(vector<C> f, int dir=1) {
  const int n = f.size();
  if (n < 2) return f;
  vector<C> f0(n/2), f1(n/2);
  rep (i, n/2) f0[i] = f[i*2], f1[i] = f[i*2 + 1];
  f0 = fft(f0, dir);
  f1 = fft(f1, dir);
  Real theta = 2 * PI / n;
  C z(cos(theta), dir * sin(theta)),
    ac = C(1, 0);
  rep (i, n) {
    f[i] = f0[i%(n/2)] + ac * f1[i%(n/2)];
    ac *= z;
  }
  return f;
}

vector<C>
convolution(vector<C> f, vector<C> g) {
  int n = std::max<unsigned>(f.size(), g.size());
  int m; for (m = 1; m < n; m <<= 1); m <<= 1;
  f.resize(m); g.resize(m);
  f = fft(f);
  g = fft(g);
  rep (i, m) f[i] *= g[i];
  auto fg = fft(f, -1);
  rep (i, m) fg[i] /= static_cast<Real>(m);
  return fg;
}

int main() {
  int n; cin >> n;
  vector<C> f(n+1), g(n+1);

  rep (i, n) {
    int a, b; cin >> a >> b;
    f[i+1] = C(a, 0);
    g[i+1] = C(b, 0);
  }

  vector<C> fg = convolution(f, g);
  const int m = fg.size();

  rep (i, 2 * n) {
    int a = round(fg[i+1].real());
    cout << a << endl;
  }

  return 0;
}