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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define int long long

const int MOD = 200003;
const int P = MOD;
const int G = 2;

namespace Fft {
typedef double ld;
typedef complex<ld> cd;

void dft(vector<cd> &a, int n) {
  vector<int> rev(1 << n);
  for (int i = 1; i < (1 << n); i++) {
    rev[i] = ((i & 1) << (n - 1));
    rev[i] |= (rev[i >> 1] >> 1);
  }
  for (int i = 0; i < (1 << n); i++) {
    if (i < rev[i]) {
      swap(a[i], a[rev[i]]);
    }
  }
  for (int step = 0; step < n; step++) {
    ld alpha = M_PI / (1 << step);
    cd w(cos(alpha), sin(alpha));
    for (int i = 0; i < (1 << n); i += (1 << (step + 1))) {
      cd curw(1, 0);
      for (int j = 0; j < (1 << step); j++) {
        cd u = a[i + j];
        cd v = curw * a[i + j + (1 << step)];
        a[i + j] = u + v;
        a[i + j + (1 << step)] = u - v;
        curw *= w;
      }
    }
  }
}

vector<int> mul(const vector<int> &a, const vector<int> &b) {
  int to = (int)max(a.size(), b.size());
  while (__builtin_popcount(to) > 1 || to < (int)(a.size() + b.size())) {
    to++;
  }
  int p = 0;
  while ((1 << p) < to) {
    p++;
  }
  vector<cd> f(to);
  for (int i = 0; i < to; i++) {
    if (i < (int)a.size()) {
      f[i] += cd(a[i], 0);
    }
    if (i < (int)b.size()) {
      f[i] += cd(0, b[i]);
    }
  }
  dft(f, p);
  for (int i = 0; i < to; i++) {
    f[i] *= f[i];
  }
  dft(f, p);
  reverse(f.begin() + 1, f.end());
  vector<int> ans;
  for (auto &t : f) {
    t /= to;
    ans.push_back(round(t.imag() / 2));
  }
  return ans;
}
} // namespace Fft

signed main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

  int n;
  cin >> n;
  vector<int> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  vector<int> wr(P, -1);
  vector<int> pw(P);
  {
    int c = 1;
    for (int i = 0; i < P; i++) {
      wr[c] = i;
      pw[i] = c;
      c = (c * G) % P;
    }
  }
  vector<int> cnt(P);
  for (int i = 0; i < n; i++) {
    if (a[i] > 0) {
      cnt[wr[a[i]]]++;
    }
  }
  vector<int> res = Fft::mul(cnt, cnt);
  int ans = 0;
  {
    for (int i = 0; i < 2 * P; i++) {
      ans += (pw[i % (P - 1)] * res[i]);
    }
  }
  for (int i = 0; i < n; i++) {
    ans -= (a[i] * a[i]) % MOD;
  }
  ans /= 2;
  cout << ans << '\\n';
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

#define int long long

const int MOD = 200003;
const int P = MOD;

signed main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);

  int n;
  cin >> n;
  vector<int> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  int F = 0;
  {
    int c = 0;
    for (int i = n - 1; i >= 0; i--) {
      F += c * a[i];
      c += a[i];
    }
  }
  int S = 0;
  {
    vector<int> na;
    for (auto t : a) {
      if (t > 0) {
        na.push_back(t);
      }
    }
    a = na;
    n = (int)a.size();
    sort(a.begin(), a.end());
    for (int c = 1; c < P; c++) {
      int pnt = 0;
      while (pnt < n) {
        int val = (c * P + a[pnt] - 1) / a[pnt];
        int L = 0;
        int R = n - 1 - pnt + 1;
        while (R - L > 1) {
          int M = (R + L) >> 1;
          int mid_val = (c * P + a[pnt + M] - 1) / a[pnt + M];
          if (val == mid_val) {
            L = M;
          } else {
            R = M;
          }
        }
        int cnt = a.end() - lower_bound(a.begin(), a.end(), val);
        S += (L + 1) * (cnt);
        pnt += R;
      }
      // for (int i = 0; i < n; i++) {
      // 	int val = (c * P + a[i] - 1) / a[i];
      // 	S += a.end() - lower_bound(a.begin(), a.end(), val);
      // }
    }
    for (int i = 0; i < n; i++) {
      S -= ((a[i] * a[i]) / MOD);
    }
    S /= 2;
  }
  cout << F - S * MOD << endl;
}