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

#include <climits>
#include <cmath>
#include <complex>
#include <iostream>
#include <map>
#include <vector>
using namespace std;

typedef pair<int, int> i_i;
typedef double D;
typedef complex<D> P;
typedef long long ll;
D pi = acos(-1.0);

void FFT(vector<P> &a, int N, double t) {
  for (int m = N; m >= 2; m /= 2) {
    int mh = m / 2;
    for (int i = 0; i < mh; i++) {
      P w = polar(1.0, t * i);
      for (int j = i; j < N; j += m) {
        int k = j + mh;
        P x = a[j] - a[k];
        a[j] += a[k];
        a[k] = w * x;
      }
    }
    t *= 2;
  }
  int i = 0;
  for (int j = 1; j <= N - 2; j++) {
    for (int k = N / 2; k > (i ^= k); k /= 2)
      ;
    if (j < i)
      swap(a[i], a[j]);
  }
}

vector<P> conv(vector<P> a, vector<P> b) {
  int n = a.size() + b.size();
  int N = 1;
  while (N < n)
    N *= 2;
  a.resize(N);
  b.resize(N);
  FFT(a, N, 2 * pi / N);
  FFT(b, N, 2 * pi / N);
  for (int i = 0; i < N; i++)
    a[i] *= b[i];
  FFT(a, N, -2 * pi / N);
  for (int i = 0; i < N; i++)
    a[i] /= N;
  return a;
}

vector<int> f(int N, vector<int> &a) {
  vector<int> sum(N + 1);
  for (int i = 0; i < N; i++)
    sum[i + 1] = sum[i] + a[i];
  vector<P> p(200001), q = p;
  for (int i = 0; i <= N; i++) {
    p[sum[i]] += 1;
    q[200000 - sum[i]] += 1;
  }
  vector<P> r = conv(p, q);
  vector<int> ans(200001);
  for (int x = 1; x <= 200000; x++)
    ans[x] += (int)(r[200000 + x].real() + 0.5);
  return ans;
}

int main() {
  int N, M, Q;
  cin >> N >> M >> Q;
  vector<int> a(N), b(M);
  for (int i = 0; i < N; i++)
    scanf("%d", &a[i]);
  for (int j = 0; j < M; j++)
    scanf("%d", &b[j]);
  vector<int> A = f(N, a), B = f(M, b);
  vector<P> p(200001), q(200001);
  for (int x = 0; x <= 200000; x++)
    p[x] += A[x];
  for (int x = 0; x <= 200000; x++)
    q[200000 - x] += B[x];
  vector<P> r = conv(p, q);
  vector<ll> ans(200001);
  for (int x = 0; x <= 200000; x++)
    ans[x] +=
        (ll)(r[200000 + x].real() + 0.5) + (ll)(r[200000 - x].real() + 0.5);
  ans[0] /= 2;
  while (Q--) {
    int c;
    scanf("%d", &c);
    printf("%lld\\n", ans[c]);
  }
}

Error: TLE

Fixed submission:

#include <climits>
#include <cmath>
#include <complex>
#include <iostream>
#include <map>
#include <vector>
using namespace std;

typedef pair<int, int> i_i;
typedef long double D;
typedef complex<D> P;
typedef long long ll;
D pi = acos(-1.0);

void FFT(vector<P> &a, int N, double t) {
  for (int m = N; m >= 2; m /= 2) {
    int mh = m / 2;
    for (int i = 0; i < mh; i++) {
      P w = polar(1.0, t * i);
      for (int j = i; j < N; j += m) {
        int k = j + mh;
        P x = a[j] - a[k];
        a[j] += a[k];
        a[k] = w * x;
      }
    }
    t *= 2;
  }
  int i = 0;
  for (int j = 1; j <= N - 2; j++) {
    for (int k = N / 2; k > (i ^= k); k /= 2)
      ;
    if (j < i)
      swap(a[i], a[j]);
  }
}

vector<P> conv(vector<P> a, vector<P> b) {
  int n = a.size() + b.size();
  int N = 1;
  while (N < n)
    N *= 2;
  a.resize(N);
  b.resize(N);
  FFT(a, N, 2 * pi / N);
  FFT(b, N, 2 * pi / N);
  for (int i = 0; i < N; i++)
    a[i] *= b[i];
  FFT(a, N, -2 * pi / N);
  for (int i = 0; i < N; i++)
    a[i] /= N;
  return a;
}

vector<int> f(int N, vector<int> &a) {
  vector<int> sum(N + 1);
  for (int i = 0; i < N; i++)
    sum[i + 1] = sum[i] + a[i];
  vector<P> p(200001), q = p;
  for (int i = 0; i <= N; i++) {
    p[sum[i]] += 1;
    q[200000 - sum[i]] += 1;
  }
  vector<P> r = conv(p, q);
  vector<int> ans(200001);
  for (int x = 1; x <= 200000; x++)
    ans[x] += (int)(r[200000 + x].real() + 0.5);
  return ans;
}

int main() {
  int N, M, Q;
  cin >> N >> M >> Q;
  vector<int> a(N), b(M);
  for (int i = 0; i < N; i++)
    scanf("%d", &a[i]);
  for (int j = 0; j < M; j++)
    scanf("%d", &b[j]);
  vector<int> A = f(N, a), B = f(M, b);
  vector<P> p(200001), q(200001);
  for (int x = 0; x <= 200000; x++)
    p[x] += A[x];
  for (int x = 0; x <= 200000; x++)
    q[200000 - x] += B[x];
  vector<P> r = conv(p, q);
  vector<ll> ans(200001);
  for (int x = 0; x <= 200000; x++)
    ans[x] +=
        (ll)(r[200000 + x].real() + 0.5) + (ll)(r[200000 - x].real() + 0.5);
  ans[0] /= 2;
  while (Q--) {
    int c;
    scanf("%d", &c);
    printf("%lld\\n", ans[c]);
  }
}