Problem p01069 (C++) — buggy submission, verdict: Runtime Error

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;
using ld = long double;
const ld eps = 1e-9;
namespace FastFourierTransform {
using C = complex<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 w, s, t;
  for (int i = 1; i < N; i <<= 1) {
    for (int k = 0; k < i; k++) {
      w = polar(1.0, PI / i * k);
      for (int j = 0; j < N; j += i * 2) {
        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;
      }
    }
  }
  if (rev)
    for (int i = 0; i < N; i++)
      F[i] /= N;
}

vector<long long> Multiply(const vector<long long int> &A,
                           const vector<long long int> &B) {
  int sz = 1;
  while (sz <= A.size() + B.size())
    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);
  vector<long long> X(A.size() + B.size() - 1);
  for (int i = 0; i < A.size() + B.size() - 1; i++)
    X[i] = F[i].real() + 0.5;
  return (X);
}
}; // namespace FastFourierTransform
const int asize = 2e5 + 100;
vector<long long int> get_mp(const vector<int> v) {
  vector<long long int> sums(v.size() + 1);
  for (int i = 0; i < v.size(); ++i) {
    sums[i + 1] = sums[i] + v[i];
  }

  vector<long long int> a(asize);
  vector<long long int> b(asize);
  for (auto sum : sums) {
    a[sum]++;
    b[asize - sum - 1]++;
  }
  vector<long long int> c = FastFourierTransform::Multiply(a, b);
  vector<long long int> ans(asize);
  for (int i = asize; i < 2 * asize - 1; ++i) {
    if (c[i]) {
      ans[i - asize + 1] += c[i];
    }
  }
  return ans;
}
vector<long long int> solve(vector<long long int> a, vector<long long int> b) {
  reverse(a.begin(), a.end());

  vector<long long> c = FastFourierTransform::Multiply(a, b);
  vector<long long int> ans(asize);
  ans[0] = c[asize - 1];
  for (int i = 1; i < asize; ++i) {
    ans[i] = c[asize - 1 - i] + c[asize - 1 + i];
  }
  return ans;
}

int main() {
  int N, M, Q;
  cin >> N >> M >> Q;
  vector<int> as(N), bs(M);
  for (int i = 0; i < N; ++i) {
    cin >> as[i];
  }
  for (int i = 0; i < M; ++i) {
    cin >> bs[i];
  }
  long long int aaa = static_cast<long long int>(N) * (N + 1) / 2;
  long long int bbb = static_cast<long long int>(M) * (M + 1) / 2;
  auto a_mp(get_mp(as));
  assert(aaa = accumulate(a_mp.begin(), a_mp.end(), 0ll));
  auto b_mp(get_mp(bs));
  assert(bbb = accumulate(b_mp.begin(), b_mp.end(), 0ll));
  auto ans(solve(a_mp, b_mp));
  assert(abs(aaa * bbb - accumulate(ans.begin(), ans.end(), 0ll)) < 1e9);
  while (Q--) {
    int n;
    cin >> n;
    cout << ans[n] << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;
using ld = long double;
const ld eps = 1e-9;
/*
?????????????????¨?????? FFT ???????????????
n?¬??????????g(x) = \\sum_{i=0}^n a_i x^i , h(x) = \\sum_{i=0}^n b_i x^i
??????????????? f(x) = (g*h)(x) = \\sum_{k=0}^{2n} \\sum_{i=0}^{k} a_i b_{k-i} x^k
???O(n log n)??§?±??????????

?????????:
g(x)????????°??? vector<T> g
h(x)????????°??? vector<T> h
??¨????????¨???
f = multiply(g, h);
??§f(x)????????°vector<T> f????±??????????

verified by: http://atc001.contest.atcoder.jp/tasks/fft_c
(AtCoder Typical Contest 001 C)

complex, vector, algorithm ?????????????????????????????????????????????
*/
#include <algorithm>
#include <complex>
#include <vector>
using namespace std;
const double PI = 3.141592653589793;
/*
 *
 * n?¬??????????f????????????????????¨????????????
 * is_inv = false => ?????????
 * is_inv = true  => ?????????
 * ?????????????????¨??????multiply????????§n??§?????£?????????
 */
using comp = complex<double>;
void fft(vector<comp> &f, int n, bool is_inv) {
  if (n == 1)
    return;
  if (n == 2) {
    comp zeta = comp(-1, 0);
    comp pow_zeta = 1.0;

    comp f0(f[0]);
    comp f1(f[1]);

    f[0] = f0 + f1;
    f[1] = f0 - f1;
    return;
  }
  vector<comp> f0(n / 2), f1(n / 2);
  for (int i = 0; i < n / 2; i++) {
    f0[i] = f[2 * i];
    f1[i] = f[2 * i + 1];
  }
  fft(f0, n / 2, is_inv);
  fft(f1, n / 2, is_inv);
  comp zeta = comp(cos(2 * PI / n), sin(2 * PI / n));
  if (is_inv)
    zeta = 1.0 / zeta;
  comp pow_zeta = 1.0;
  for (int i = 0; i < n; i++) {
    f[i] = f0[i % (n / 2)] + pow_zeta * f1[i % (n / 2)];
    pow_zeta *= zeta;
  }
}

template <typename T> vector<T> multiply(vector<T> &g, vector<T> &h) {
  int n = 1;
  int gs = g.size();
  int hs = h.size();
  int fs = gs + hs - 1;
  while (n <= fs + 1)
    n *= 2;
  vector<comp> gg(n), hh(n), ff(n);
  for (int i = 0; i < n; i++) {
    gg[i] = (i < gs) ? g[i] : 0;
    hh[i] = (i < hs) ? h[i] : 0;
  }
  fft(gg, n, false);
  fft(hh, n, false);
  for (int i = 0; i < n; i++)
    ff[i] = gg[i] * hh[i];
  fft(ff, n, true);

  vector<T> f(fs);
  for (int i = 0; i < fs; i++) {
    f[i] = (T)round(ff[i].real() / n);
  }

  return f;
}
const int asize = 2e5 + 1;
vector<long long int> get_mp(const vector<int> v) {
  vector<long long int> sums(v.size() + 1);
  for (int i = 0; i < v.size(); ++i) {
    sums[i + 1] = sums[i] + v[i];
  }

  vector<long long int> a(asize);
  vector<long long int> b(asize);
  for (auto sum : sums) {
    a[sum]++;
    b[asize - sum - 1]++;
  }
  vector<long long int> c = multiply(a, b);
  vector<long long int> ans(asize);
  for (int i = asize; i < 2 * asize - 1; ++i) {
    if (c[i]) {
      ans[i - asize + 1] += c[i];
    }
  }
  return ans;
}
vector<long long int> solve(vector<long long int> a, vector<long long int> b) {
  reverse(a.begin(), a.end());

  vector<long long> c = multiply<long long int>(a, b);
  vector<long long int> ans(asize);
  ans[0] = c[asize - 1];
  for (int i = 1; i < asize; ++i) {
    ans[i] = c[asize - 1 - i] + c[asize - 1 + i];
  }
  return ans;
}

int main() {
  int N, M, Q;
  cin >> N >> M >> Q;
  vector<int> as(N), bs(M);
  for (int i = 0; i < N; ++i) {
    cin >> as[i];
  }
  for (int i = 0; i < M; ++i) {
    cin >> bs[i];
  }
  long long int aaa = static_cast<long long int>(N) * (N + 1) / 2;
  long long int bbb = static_cast<long long int>(M) * (M + 1) / 2;
  auto a_mp(get_mp(as));
  assert(aaa = accumulate(a_mp.begin(), a_mp.end(), 0ll));
  auto b_mp(get_mp(bs));
  assert(bbb = accumulate(b_mp.begin(), b_mp.end(), 0ll));
  auto ans(solve(a_mp, b_mp));
  assert(abs(aaa * bbb - accumulate(ans.begin(), ans.end(), 0ll)) < 1e6);
  while (Q--) {
    int n;
    cin >> n;
    cout << ans[n] << endl;
  }
  return 0;
}