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

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

struct FFT {
  using CP = complex<double>;
  const double PI = acos(-1);
  int base = 1;
  vector<CP> roots{CP(0, 0), CP(1, 0)};
  vector<int> rv{0, 1};
  FFT() {}
  void ensure_base(int nb) {
    if (nb <= base)
      return;
    rv.resize(1 << nb);
    roots.resize(1 << nb);
    for (int i = 0; i < (1 << nb); ++i)
      rv[i] = (rv[i >> 1] >> 1) + ((i & 1) << (nb - 1));
    while (base < nb) {
      double arg = PI * 2.0 / (1 << (base + 1));
      for (int i = 1 << (base - 1); i < (1 << base); ++i) {
        roots[i << 1] = roots[i];
        double narg = arg * (2 * i + 1 - (1 << base));
        roots[(i << 1) + 1] = CP(cos(narg), sin(narg));
      }
      ++base;
    }
  }
  void fft(vector<CP> &a, int n, bool sg = 0) {
    assert((n & (n - 1)) == 0);
    for (int i = 0; i < n; ++i)
      if (i < rv[i])
        swap(a[i], a[rv[i]]);
    for (int k = 1; k < n; k <<= 1)
      for (int i = 0; i < n; i += 2 * k)
        for (int j = 0; j < k; ++j) {
          CP z = a[i + j + k] * (sg ? roots[j + k] : conj(roots[j + k]));
          a[i + j + k] = a[i + j] - z;
          a[i + j] = a[i + j] + z;
        }
    if (sg)
      for (int i = 0; i < n; ++i)
        a[i] /= n;
  }
  vector<long long> mul(const vector<int> &a, const vector<int> &b) {
    int need = a.size() + b.size() - 1;
    int nb = 1;
    while ((1 << nb) < need)
      ++nb;
    ensure_base(nb);
    int sz = 1 << nb;
    vector<CP> fa(sz), fb(sz);
    for (int i = 0; i < sz; ++i) {
      if (i < a.size())
        fa[i] = CP(a[i], 0);
      if (i < b.size())
        fb[i] = CP(b[i], 0);
    }
    fft(fa, sz);
    fft(fb, sz);
    for (int i = 0; i < sz; ++i)
      fa[i] *= fb[i];
    fft(fa, sz, 1);
    vector<long long> res(need);
    for (int i = 0; i < need; ++i)
      res[i] = llround(fa[i].real());
    return res;
  }
};

const int p = 200003;
int n;
vector<long long> cnt;

long long solve();

int main() {
  cin >> n;
  cnt.resize(p);
  for (int i = 0; i < n; ++i) {
    int a;
    cin >> a;
    ++cnt[a];
  }
  cout << solve() << endl;
  return 0;
}

long long solve() {
  vector<int> memo(1, 1), tmp(p);
  int now = 2;
  for (int i = 1; i < p; ++i) {
    tmp[i] = cnt[now];
    memo.push_back(now);
    (now <<= 1) %= p;
  }
  for (int i = 0; i < p; ++i) {
    memo.push_back(now);
    (now <<= 1) %= p;
  }
  FFT fft;
  auto v = fft.mul(tmp, tmp);
  long long res = 0;
  for (long long i = 0; i < 2 * p; ++i) {
    res += v[i] * memo[i];
    if (i < p)
      res -= cnt[i] * (i * i % p);
  }
  return res / 2;
}

Error: TLE

Fixed submission:

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

const long long p = 200003;
long long n;
vector<long long> cnt, sum, csum;

long long solve();

int main() {
  cin >> n;
  cnt.assign(p + 1, 0);
  csum.assign(p + 1, 0);
  sum.assign(p + 1, 0);
  for (int i = 0; i < n; ++i) {
    int a;
    cin >> a;
    ++cnt[a];
    ++csum[a];
    sum[a] += a;
  }
  cout << solve() << endl;
  return 0;
}

long long solve() {
  long long res = 0;
  for (int i = 1; i <= p; ++i)
    sum[i] += sum[i - 1];
  for (int i = 1; i <= p; ++i)
    csum[i] += csum[i - 1];
  for (long long i = 1; i <= p; ++i)
    if (cnt[i]) {
      long long now = i * i % p, id = 0;
      res += cnt[i] * (cnt[i] - 1) / 2 * now;
      while (id < i) {
        now = id * i / p;
        long long r = ((now + 1) * p + i - 1) / i, tmp = 0;
        r = min(r, i);
        if (r > 0)
          tmp = sum[r - 1];
        if (id > 0)
          tmp -= sum[id - 1];
        long long aa = tmp * i * cnt[i];
        res += tmp * i * cnt[i];
        tmp = 0;
        if (r > 0)
          tmp = csum[r - 1];
        if (id > 0)
          tmp -= csum[id - 1];
        aa -= tmp * now * cnt[i] * p;
        res -= tmp * now * cnt[i] * p;

        id = r;
      }
    }
  return res;
}