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

#include <bits/stdc++.h>
using namespace std;

#define int long long
typedef pair<int, int> pint;
typedef vector<int> vint;
typedef vector<pint> vpint;
#define pb push_back
#define fi first
#define se second
#define all(v) (v).begin(), (v).end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define each(it, v)                                                            \\
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
template <class T, class U> inline void chmin(T &t, U f) {
  if (t > f)
    t = f;
}
template <class T, class U> inline void chmax(T &t, U f) {
  if (t < f)
    t = f;
}

const int mod[3] = {998244353, 897581057, 645922817};
struct NumberTheoreticTransform {
  int mod;
  int root = 3;
  NumberTheoreticTransform(int mod) : mod(mod) {}

  int mul(int x, int y) { return x * y % mod; }
  int add(int x, int y) { return (x += y) >= mod ? x - mod : x; }

  int pow(int x, int y) {
    int res = 1;
    while (y) {
      if (y & 1)
        res = mul(res, x);
      x = mul(x, x);
      y >>= 1;
    }
    return res;
  }

  int inv(int x) { return pow(x, mod - 2); }

  void ntt(vector<int> &f, bool rev = false) {
    int n = f.size();
    for (int i = 0, j = 1; j < n - 1; j++) {
      for (int k = n >> 1; k > (i ^= k); k /= 2)
        ;
      if (i > j)
        swap(f[i], f[j]);
    }

    for (int m = 2; m <= n; m *= 2) {
      int wr = pow(root, (mod - 1) / m);
      if (rev)
        wr = inv(wr);
      for (int i = 0; i < n; i += m) {
        int w = 1;
        for (int j = 0; j < m / 2; j++) {
          int f0 = f[i + j], f1 = mul(w, f[i + j + m / 2]);
          f[i + j] = add(f0, f1);
          f[i + j + m / 2] = add(f0, mod - f1);
          w = mul(w, wr);
        }
      }
    }
    if (rev) {
      int v = inv(n);
      for (int i = 0; i < n; i++)
        f[i] = mul(f[i], v);
    }
  }
  vector<int> multiply(vector<int> A, vector<int> B) {
    int n = 1;
    while (n < A.size() + B.size() - 1)
      n <<= 1;
    A.resize(n);
    B.resize(n);
    ntt(A);
    ntt(B);
    for (int i = 0; i < n; i++)
      A[i] = mul(A[i], B[i]);
    ntt(A, true);
    A.resize(A.size() + B.size() - 1);
    return A;
  }
};

int modPow(int n, int m, int mod) {
  int res = 1;
  while (m) {
    if (m & 1)
      res = res * n % mod;
    n = n * n % mod;
    m >>= 1;
  }
  return res;
}

inline int extgcd(int a, int b, int &x, int &y) {
  x = 1, y = 0;
  int g = a;
  if (b != 0)
    g = extgcd(b, a % b, y, x), y -= a / b * x;
  return g;
}
inline int INV(int a, int mod) {
  int x, y;
  extgcd(a, mod, x, y);
  return (x % mod + mod) % mod;
}
int chineneRemainder(const vector<int> b, const vector<int> m) {
  int n = b.size();
  vector<int> constant(n), coef(n, 1), v(n);
  for (int i = 0; i < n; i++) {
    v[i] = (b[i] % m[i] - constant[i] % m[i] + m[i]) % m[i];
    v[i] = v[i] % m[i] * INV(coef[i], m[i]) % m[i];
    for (int j = i + 1; j < n; j++) {
      constant[j] = (constant[j] + v[i] * coef[j]) % m[j];
      coef[j] = coef[j] * m[i] % m[j];
    }
  }
  int res = 0;
  for (int i = n - 1; i >= 0; i--)
    res = res * m[i] + v[i];
  return res;
}

const int lim = 200000;
int n, m, q;
vector<int> a, b;

int sum[222222];
vector<int> get(vector<int> &a, int n) {
  rep(i, a.size()) sum[i + 1] = sum[i] + a[i];

  vector<int> g(lim + 1), h(lim + 1);
  rep(i, n) {
    g[sum[i + 1]]++;
    h[lim - sum[i]]++;
  }
  NumberTheoreticTransform ntt(mod[0]);
  vector<int> f = ntt.multiply(g, h);
  vector<int> ret(lim + 1);
  for (int i = 1; i <= lim; i++)
    ret[i] = f[i + lim];
  return ret;
}

int ans[lim + 1];

signed main() {
  scanf("%lld%lld%lld", &n, &m, &q);
  a.resize(n);
  b.resize(m);
  rep(i, n) scanf("%lld", &a[i]);
  rep(i, m) scanf("%lld", &b[i]);

  vector<int> aa = get(a, n);
  vector<int> bb = get(b, m);

  vector<vint> tmp(3, vint(lim + 1));
  for (int t = 0; t < 3; t++) {
    vector<int> cc(lim + 1);
    rep(i, lim + 1) cc[lim - i] = bb[i];
    NumberTheoreticTransform ntt(mod[t]);
    vector<int> f = ntt.multiply(aa, cc);
    for (int i = -lim; i <= lim; i++)
      tmp[t][abs(i)] += f[lim + i];
  }

  for (int i = 0; i <= lim; i++) {
    vector<int> b = {tmp[0][i], tmp[1][i], tmp[2][i]};
    vector<int> m = {mod[0], mod[1], mod[2]};
    ans[i] = chineneRemainder(b, m);
  }

  while (q--) {
    int c;
    scanf("%lld", &c);
    printf("%lld\\n", ans[c]);
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define int long long
typedef pair<int, int> pint;
typedef vector<int> vint;
typedef vector<pint> vpint;
#define pb push_back
#define fi first
#define se second
#define all(v) (v).begin(), (v).end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define each(it, v)                                                            \\
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
template <class T, class U> inline void chmin(T &t, U f) {
  if (t > f)
    t = f;
}
template <class T, class U> inline void chmax(T &t, U f) {
  if (t < f)
    t = f;
}

namespace FFT {
typedef complex<long double> C;
const long double PI = acos(-1);

void dft(vector<C> &f, bool rev = false) {
  int n = f.size();
  for (int i = 0, j = 1; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k /= 2)
      ;
    if (i > j)
      swap(f[i], f[j]);
  }
  for (int m = 2; m <= n; m *= 2) {
    C wr = polar((long double)1.0, (rev ? -1 : 1) * 2.0 * PI / m);
    for (int i = 0; i < n; i += m) {
      C w(1, 0);
      for (int j = 0; j < m / 2; j++) {
        C f0 = f[i + j], f1 = w * f[i + j + m / 2];
        f[i + j] = f0 + f1;
        f[i + j + m / 2] = f0 - f1;
        w *= wr;
      }
    }
  }
  if (rev)
    for (int i = 0; i < n; i++)
      f[i] *= 1.0 / n;
}

vector<int> multiply(vector<int> A, vector<int> B) {
  int n = 1;
  while (n <= A.size() + B.size())
    n <<= 1;
  vector<C> g(n), h(n);
  for (int i = 0; i < A.size(); i++)
    g[i] = A[i];
  for (int i = 0; i < B.size(); i++)
    h[i] = B[i];

  dft(g);
  dft(h);
  for (int i = 0; i < n; i++)
    g[i] *= h[i];
  dft(g, true);

  vector<int> C(A.size() + B.size() - 1);
  for (int i = 0; i < C.size(); i++)
    C[i] = (int)(g[i].real() + 0.5);
  return C;
}
} // namespace FFT

int n, m, q;
vector<int> a, b;

int sum[222222];
vector<int> get(vector<int> &a, int n) {
  rep(i, a.size()) sum[i + 1] = sum[i] + a[i];

  vector<int> g(200001), h(200001);
  rep(i, n) {
    g[sum[i + 1]]++;
    h[200000 - sum[i]]++;
  }
  vector<int> f = FFT::multiply(g, h);
  vector<int> ret(200001);
  for (int i = 1; i <= 200000; i++)
    ret[i] = f[i + 200000];
  return ret;
}

int ans[222222];

signed main() {
  scanf("%lld%lld%lld", &n, &m, &q);
  a.resize(n);
  b.resize(m);
  rep(i, n) scanf("%lld", &a[i]);
  rep(i, m) scanf("%lld", &b[i]);

  vector<int> aa = get(a, n);
  vector<int> bb = get(b, m);

  vector<int> cc(200001);
  rep(i, 200001) cc[200000 - i] = bb[i];

  vector<int> f = FFT::multiply(aa, cc);
  for (int i = -200000; i <= 200000; i++)
    ans[abs(i)] += f[200000 + i];

  while (q--) {
    int c;
    scanf("%lld", &c);
    printf("%lld\\n", ans[c]);
  }
  return 0;
}