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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD @
#define ADD(X, Y) ((X) = ((X) + (Y) % MOD) % MOD)
typedef long long i64;
typedef vector<int> ivec;
typedef vector<string> svec;

const double PI = acos(-1);
struct CP {
  double x, y;
  CP(double x = 0, double y = 0) : x(x), y(y) {}
  CP operator+(const CP &a) { return CP(x + a.x, y + a.y); }
  CP operator-(const CP &a) { return CP(x - a.x, y - a.y); }
  CP operator*(const CP &a) { return CP(x * a.x - y * a.y, x * a.y + y * a.x); }
};
void FFTMakeG(int n, CP g[]) {
  for (int i = 0; i < n; ++i) {
    double th = PI * 2.0 * i / n;
    g[i] = CP(cos(th), sin(th));
  }
}
void FFT(int n, CP g[], CP z[]) {
  int i, j, k, l, m;
  for (i = 1, j = 0; i < n; ++i) {
    for (k = n >> 1; (j ^= k) < k; k >>= 1)
      ;
    if (i < j)
      swap(z[i], z[j]);
  }
  for (l = k = 1, m = __builtin_ctz(n); m--; l |= k <<= 1) {
    for (i = 0; i < n; ++i)
      if (i < (j = i ^ k)) {
        CP t = z[j] * g[(i & l) << m];
        z[j] = z[i] - t;
        z[i] = z[i] + t;
      }
  }
}

const int P = 200003;
int N, A[202020];
i64 pows[202020];
map<i64, int> ord;
int cnt[202020];

CP g[1 << 19], xhi[1 << 19], xlo[1 << 19], y[1 << 19];
const int K = 330;
i64 ans[202020];

i64 naive() {
  i64 ret = 0;
  for (int i = 0; i < N; ++i) {
    for (int j = i + 1; j < N; ++j)
      ret += (A[i] * (i64)A[j]) % P;
  }
  return ret;
}

int main() {
  scanf("%d", &N);
  for (int i = 0; i < N; ++i)
    scanf("%d", A + i);

  i64 v = 1;
  for (int i = 0; i < P - 1; ++i) {
    ord[v] = i;
    pows[i] = v;
    v = v * 2 % P;
  }
  for (int i = 0; i < N; ++i) {
    int ai = A[i];
    if (ai == 0)
      continue;
    ++cnt[ord[ai]];
  }
  FFTMakeG(1 << 19, g);
  for (int i = 0; i < P - 1; ++i) {
    xhi[i] = cnt[i] / K;
    xlo[i] = cnt[i] % K;
  }
  FFT(1 << 19, g, xhi);
  FFT(1 << 19, g, xlo);

  for (int i = 0; i < 3; ++i) {
    if (i == 0) {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xhi[j] * xhi[j];
    } else if (i == 1) {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xhi[j] * xlo[j];
    } else {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xlo[j] * xlo[j];
    }
    reverse(g + 1, g + (1 << 19));
    FFT(1 << 19, g, y);
    reverse(g + 1, g + (1 << 19));

    i64 waffle = 1;
    if (i == 0)
      waffle = K * K;
    else if (i == 1)
      waffle = K * 2;

    for (int j = 0; j < (1 << 19); ++j) {
      i64 v = (i64)(y[j].x / (1 << 19) + 0.5);
      ans[j % (P - 1)] += v * waffle;
    }
  }
  i64 ret = 0;
  for (int i = 0; i < P - 1; ++i) {
    ret += ans[i] * pows[i];
    ret -= cnt[i] * pows[(i * 2) % (P - 1)];
  }
  ret /= 2;
  printf("%lld\\n", ret);
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD @
#define ADD(X, Y) ((X) = ((X) + (Y) % MOD) % MOD)
typedef long long i64;
typedef vector<int> ivec;
typedef vector<string> svec;

const double PI = acos(-1);
struct CP {
  double x, y;
  CP(double x = 0, double y = 0) : x(x), y(y) {}
  CP operator+(const CP &a) { return CP(x + a.x, y + a.y); }
  CP operator-(const CP &a) { return CP(x - a.x, y - a.y); }
  CP operator*(const CP &a) { return CP(x * a.x - y * a.y, x * a.y + y * a.x); }
};
void FFTMakeG(int n, CP g[]) {
  for (int i = 0; i < n; ++i) {
    double th = PI * 2.0 * i / n;
    g[i] = CP(cos(th), sin(th));
  }
}
void FFT(int n, CP g[], CP z[]) {
  int i, j, k, l, m;
  for (i = 1, j = 0; i < n; ++i) {
    for (k = n >> 1; (j ^= k) < k; k >>= 1)
      ;
    if (i < j)
      swap(z[i], z[j]);
  }
  for (l = k = 1, m = __builtin_ctz(n); m--; l |= k <<= 1) {
    for (i = 0; i < n; ++i)
      if (i < (j = i ^ k)) {
        CP t = z[j] * g[(i & l) << m];
        z[j] = z[i] - t;
        z[i] = z[i] + t;
      }
  }
}

const int P = 200003;
int N, A[202020];
i64 pows[202020];
map<i64, int> ord;
int cnt[202020];

CP g[1 << 19], xhi[1 << 19], xlo[1 << 19], y[1 << 19];
const int K = 330;
i64 ans[202020];

i64 naive() {
  i64 ret = 0;
  for (int i = 0; i < N; ++i) {
    for (int j = i + 1; j < N; ++j)
      ret += (A[i] * (i64)A[j]) % P;
  }
  return ret;
}

int main() {
  scanf("%d", &N);
  for (int i = 0; i < N; ++i)
    scanf("%d", A + i);

  i64 v = 1;
  for (int i = 0; i < P - 1; ++i) {
    ord[v] = i;
    pows[i] = v;
    v = v * 2 % P;
  }
  for (int i = 0; i < N; ++i) {
    int ai = A[i];
    if (ai == 0)
      continue;
    ++cnt[ord[ai]];
  }
  FFTMakeG(1 << 19, g);
  for (int i = 0; i < P - 1; ++i) {
    xhi[i] = cnt[i] / K;
    xlo[i] = cnt[i] % K;
  }
  FFT(1 << 19, g, xhi);
  FFT(1 << 19, g, xlo);

  for (int i = 0; i < 3; ++i) {
    if (i == 0) {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xhi[j] * xhi[j];
    } else if (i == 1) {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xhi[j] * xlo[j];
    } else {
      for (int j = 0; j < (1 << 19); ++j)
        y[j] = xlo[j] * xlo[j];
    }
    reverse(g + 1, g + (1 << 19));
    FFT(1 << 19, g, y);
    reverse(g + 1, g + (1 << 19));

    i64 waffle = 1;
    if (i == 0)
      waffle = K * K;
    else if (i == 1)
      waffle = K * 2;

    for (int j = 0; j < (1 << 19); ++j) {
      i64 v = (i64)(y[j].x / (1 << 19) + 0.5);
      ans[j % (P - 1)] += v * waffle;
    }
  }
  i64 ret = 0;
  for (int i = 0; i < P - 1; ++i) {
    ret += ans[i] * pows[i];
    ret -= cnt[i] * pows[(i * 2) % (P - 1)];
  }
  ret /= 2;
  if (ret != naive())
    printf("%lld\\n", naive());
  printf("%lld\\n", ret);
  return 0;
}