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

#pragma GCC optimize("O3")
// #pragma GCC target ("avx")
#include <algorithm>
#include <cmath>
#include <complex>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <string>
#include <vector>

#ifdef LOCAL
#define eprintf(...) fprintf(stderr, __VA_ARGS__)
#else
#define NDEBUG
#define eprintf(...)                                                           \\
  do {                                                                         \\
  } while (0)
#endif
#include <cassert>

using namespace std;

typedef long long LL;
typedef vector<int> VI;

#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define EACH(i, c)                                                             \\
  for (__typeof((c).begin()) i = (c).begin(), i##_end = (c).end();             \\
       i != i##_end; ++i)

template <class T> inline void amin(T &x, const T &y) {
  if (y < x)
    x = y;
}
template <class T> inline void amax(T &x, const T &y) {
  if (x < y)
    x = y;
}
template <class Iter> void rprintf(const char *fmt, Iter begin, Iter end) {
  for (bool sp = 0; begin != end; ++begin) {
    if (sp)
      putchar(' ');
    else
      sp = true;
    printf(fmt, *begin);
  }
  putchar('\\n');
}
typedef complex<double> Complex;
const double PI = acos(-1);
const Complex I(0, 1);

// Stockham algorithm
// Hint: vector<Complex> v; fft(v.data(), v.size(), false);
void fft(Complex *x, const int n, const bool idft = false) {
  static Complex *y = NULL;
  static int y_len = 0;
  if (y_len < n) {
    if (y)
      delete[] y;
    y_len = n;
    y = new Complex[y_len];
  }
  if (idft) {
    for (int p = 0; p < n; p++)
      x[p] = conj(x[p]);
  }
  int nn = n, s = 1, eo = 0;
  for (; nn > 1;) {
    const int m = nn / 2;
    const double theta = 2.0 * PI / nn;
    for (int p = 0; p < m; p++) {
      const Complex wp(cos(p * theta), -sin(p * theta));
      for (int q = 0; q < s; q++) {
        const Complex a = x[q + s * p];
        const Complex b = x[q + s * (p + m)];
        const Complex c = a - b;
        y[q + s * (p + p)] = a + b;
        y[q + s * (p + p + 1)] =
            Complex(c.real() * wp.real() - c.imag() * wp.imag(),
                    c.real() * wp.imag() + c.imag() * wp.real());
        // y[q + s * (p + p + 1)] = (a - b) * wp;
      }
    }
    nn = m;
    s += s;
    eo ^= 1;
    swap(x, y);
  }
  if (eo) {
    for (int q = 0; q < s; q++)
      y[q] = x[q];
    swap(x, y);
  }
  if (idft) {
    double b = 1.0 / n;
    for (int p = 0; p < n; p++)
      x[p] = conj(x[p]) * b;
  }
}

const LL MOD = 200003;
const int LEN = 1 << 19;

int N;
int C[MOD + 11];

LL A[MOD];
Complex X[LEN];

void MAIN() {
  scanf("%d", &N);
  REP(i, N) { scanf("%lld", A + i); }

  sort(A, A + N);

  REP(i, N) C[A[i]]++;

  {
    LL x = 1;
    int e = 0;
    // 2^e == x;
    do {
      X[e].real(C[x]);
      x = x * 2 % MOD;
      e++;
    } while (x != 1);
  }

  fft(X, LEN);
  REP(i, LEN) X[i] *= X[i];
  fft(X, LEN, true);

  LL ans = 0;
  for (LL x = 1, e = 0; e < LEN; e++, x = x * 2 % MOD) {
    ans += llround(X[e].real()) * x;
  }

  REP(i, N) ans -= A[i] * A[i] % MOD;
  assert(ans % 2 == 0);
  ans /= 2;
  printf("%lld\\n", ans);
}

int main() {
  int TC = 1;
  //    scanf("%d", &TC);
  REP(tc, TC) MAIN();
  return 0;
}


Error: TLE

Fixed submission:

#pragma GCC optimize("O3")
// #pragma GCC target ("avx")
#include <algorithm>
#include <cmath>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <string>
#include <vector>

#ifdef LOCAL
#define eprintf(...) fprintf(stderr, __VA_ARGS__)
#else
#define NDEBUG
#define eprintf(...)                                                           \\
  do {                                                                         \\
  } while (0)
#endif
#include <cassert>

using namespace std;

typedef long long LL;
typedef vector<int> VI;

#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define EACH(i, c)                                                             \\
  for (__typeof((c).begin()) i = (c).begin(), i##_end = (c).end();             \\
       i != i##_end; ++i)

template <class T> inline void amin(T &x, const T &y) {
  if (y < x)
    x = y;
}
template <class T> inline void amax(T &x, const T &y) {
  if (x < y)
    x = y;
}
template <class Iter> void rprintf(const char *fmt, Iter begin, Iter end) {
  for (bool sp = 0; begin != end; ++begin) {
    if (sp)
      putchar(' ');
    else
      sp = true;
    printf(fmt, *begin);
  }
  putchar('\\n');
}

const LL MOD = 200003;
int N;
int C[MOD + 11];
LL sums[MOD + 11];

LL A[MOD];

void MAIN() {
  scanf("%d", &N);
  REP(i, N) { scanf("%lld", A + i); }

  sort(A, A + N);

  REP(i, N) C[A[i]]++;
  REP(i, MOD) sums[i + 1] = sums[i] + C[i];

  LL neg = 0;
  for (int a = 2; a < MOD; a++)
    if (C[a]) {
      for (LL t = 1;; t++) {
        LL v = (t * MOD + a - 1) / a;
        if (v >= MOD)
          break;
        neg += (sums[MOD] - sums[v]) * C[a];
      }
    }

  LL ans = 0;
  LL sum = 0;
  REP(i, N) sum += A[i];
  REP(i, N) {
    ans += A[i] * sum;
    ans -= A[i] * A[i] % MOD;
  }

  ans -= neg * MOD;
  assert(ans % 2 == 0);
  ans /= 2;
  printf("%lld\\n", ans);
}

int main() {
  int TC = 1;
  //    scanf("%d", &TC);
  REP(tc, TC) MAIN();
  return 0;
}