Problem p01069 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;
using Int = long long;

namespace FFT {
using dbl = double;

struct num {
  dbl x, y;

  num() { x = y = 0; }

  num(dbl x, dbl y) : x(x), y(y) {}
};

inline num operator+(num a, num b) { return num(a.x + b.x, a.y + b.y); }

inline num operator-(num a, num b) { return num(a.x - b.x, a.y - b.y); }

inline num operator*(num a, num b) {
  return num(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
}

inline num conj(num a) { return num(a.x, -a.y); }

int base = 1;
vector<num> rts = {{0, 0}, {1, 0}};
vector<int> rev = {0, 1};

const dbl PI = acosl(-1.0);

void ensure_base(int nbase) {
  if (nbase <= base)
    return;

  rev.resize(1 << nbase);
  for (int i = 0; i < (1 << nbase); i++)
    rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (nbase - 1));

  rts.resize(1 << nbase);
  while (base < nbase) {
    dbl angle = 2 * PI / (1 << (base + 1));
    for (int i = 1 << (base - 1); i < (1 << base); i++) {
      rts[i << 1] = rts[i];
      dbl angle_i = angle * (2 * i + 1 - (1 << base));
      rts[(i << 1) + 1] = num(cos(angle_i), sin(angle_i));
    }
    base++;
  }
}

void fft(vector<num> &a, int n = -1) {
  if (n == -1)
    n = a.size();
  assert((n & (n - 1)) == 0);

  int zeros = __builtin_ctz(n);
  ensure_base(zeros);
  int shift = base - zeros;
  for (int i = 0; i < n; i++)
    if (i < (rev[i] >> shift))
      swap(a[i], a[rev[i] >> shift]);

  for (int k = 1; k < n; k <<= 1) {
    for (int i = 0; i < n; i += 2 * k) {
      for (int j = 0; j < k; j++) {
        num z = a[i + j + k] * rts[j + k];
        a[i + j + k] = a[i + j] - z;
        a[i + j] = a[i + j] + z;
      }
    }
  }
}

vector<num> fa;

vector<Int> multiply(vector<int> &a, vector<int> &b) {
  int need = a.size() + b.size() - 1;
  int nbase = 0;
  while ((1 << nbase) < need)
    nbase++;
  ensure_base(nbase);

  int sz = 1 << nbase;
  if (sz > (int)fa.size())
    fa.resize(sz);
  for (int i = 0; i < sz; i++) {
    int x = (i < (int)a.size() ? a[i] : 0);
    int y = (i < (int)b.size() ? b[i] : 0);
    fa[i] = num(x, y);
  }
  fft(fa, sz);

  num r(0, -0.25 / sz);
  for (int i = 0; i <= (sz >> 1); i++) {
    int j = (sz - i) & (sz - 1);
    num z = (fa[j] * fa[j] - conj(fa[i] * fa[i])) * r;
    if (i != j)
      fa[j] = (fa[i] * fa[i] - conj(fa[j] * fa[j])) * r;
    fa[i] = z;
  }
  fft(fa, sz);

  vector<Int> res(need);
  for (int i = 0; i < need; i++)
    res[i] = fa[i].x + 0.5;

  return res;
}

}; // namespace FFT

int main() {
  int N, M, Q, A[40001], B[40001], C[100000];
  vector<int> F(200001), G(200001);

  scanf("%d %d %d", &N, &M, &Q);
  for (int i = 0; i < N; i++)
    scanf("%d", &A[i + 1]), A[i + 1] += A[i];
  for (int i = 0; i < M; i++)
    scanf("%d", &B[i + 1]), B[i + 1] += B[i];
  for (int i = 0; i < Q; i++)
    scanf("%d", &C[i]);

  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < i; j++) {
      ++F[A[i] - A[j]];
    }
  }
  for (int i = 1; i <= M; i++) {
    for (int j = 0; j < i; j++) {
      ++G[B[i] - B[j]];
    }
  }
  reverse(begin(G), end(G));
  auto X = FFT::multiply(F, G);
  for (int i = 0; i < Q; i++) {
    if (C[i] == 0)
      printf("%lld\\n", X[200000]);
    else
      printf("%lld\\n", X[200000 - C[i]] + X[200000 + C[i]]);
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

using int64 = long long;

namespace FastFourierTransform {
using C = complex<long double>;

void DiscreteFourierTransform(vector<C> &F, bool rev) {
  const int N = (int)F.size();
  const double PI = (rev ? -1 : 1) * acos(-1);
  for (int i = 0, j = 1; j + 1 < N; j++) {
    for (int k = N >> 1; k > (i ^= k); k >>= 1)
      ;
    if (i > j)
      swap(F[i], F[j]);
  }

  C wr, w, s, t;
  for (int i = 1; i < N; i <<= 1) {
    wr = polar(1.0, PI / i);
    for (int j = 0; j < N; j += i * 2) {
      w = 1.0;
      for (int k = 0; k < i; k++) {
        s = F[j + k];
        t = C(F[j + k + i].real() * w.real() - F[j + k + i].imag() * w.imag(),
              F[j + k + i].real() * w.imag() + F[j + k + i].imag() * w.real());
        F[j + k] = s + t, F[j + k + i] = s - t;
        w *= wr;
      }
    }
  }
  if (rev) {
    double temp = 1.0 / N;
    for (int i = 0; i < N; i++)
      F[i] *= temp;
  }
}

void Multiply(const vector<int64> &A, const vector<int64> &B,
              vector<int64> &X) {
  int sz = 1;
  while (sz < A.size() + B.size() - 1)
    sz <<= 1;
  vector<C> F(sz), G(sz);
  for (int i = 0; i < A.size(); i++)
    F[i] = A[i];
  for (int i = 0; i < B.size(); i++)
    G[i] = B[i];
  DiscreteFourierTransform(F, false);
  DiscreteFourierTransform(G, false);
  for (int i = 0; i < sz; i++)
    F[i] *= G[i];
  DiscreteFourierTransform(F, true);
  X.resize(A.size() + B.size() - 1);
  for (int i = 0; i < A.size() + B.size() - 1; i++)
    X[i] = F[i].real() + 0.5;
}
}; // namespace FastFourierTransform

int main() {
  int N, M, Q, A[40001], B[40001], C[100000];
  vector<int64> F(200001), G(200001);

  scanf("%d %d %d", &N, &M, &Q);
  for (int i = 0; i < N; i++)
    scanf("%d", &A[i + 1]), A[i + 1] += A[i];
  for (int i = 0; i < M; i++)
    scanf("%d", &B[i + 1]), B[i + 1] += B[i];
  for (int i = 0; i < Q; i++)
    scanf("%d", &C[i]);

  for (int i = 1; i <= N; i++) {
    for (int j = 0; j < i; j++) {
      ++F[A[i] - A[j]];
    }
  }
  for (int i = 1; i <= M; i++) {
    for (int j = 0; j < i; j++) {
      ++G[B[i] - B[j]];
    }
  }
  reverse(begin(G), end(G));
  vector<int64> X;
  FastFourierTransform::Multiply(F, G, X);
  for (int i = 0; i < Q; i++) {
    if (C[i] == 0)
      printf("%lld\\n", X[200000]);
    else
      printf("%lld\\n", X[200000 - C[i]] + X[200000 + C[i]]);
  }
}