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

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
using ll = long long;
using u64 = uint_fast64_t;
using pii = pair<int, int>;
using pll = pair<long long, long long>;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()
constexpr char ln = '\\n';
constexpr long long MOD = 1000000007;
// constexpr long long MOD = 998244353;
template <class T1, class T2> inline bool chmax(T1 &a, T2 b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T1, class T2> inline bool chmin(T1 &a, T2 b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
inline int popcount(int x) { return __builtin_popcount(x); }
inline int popcount(long long x) { return __builtin_popcountll(x); }
void print() { cout << "\\n"; }
template <class T, class... Args> void print(const T &x, const Args &...args) {
  cout << x << " ";
  print(args...);
}
///////////////////////////////////////////////////////////////////////////////////////////////////////////////

const int P = 200003;

namespace FastFourierTransform {
using real = double;

struct C {
  real x, y;

  C() : x(0), y(0) {}

  C(real x, real y) : x(x), y(y) {}

  inline C operator+(const C &c) const { return C(x + c.x, y + c.y); }

  inline C operator-(const C &c) const { return C(x - c.x, y - c.y); }

  inline C operator*(const C &c) const {
    return C(x * c.x - y * c.y, x * c.y + y * c.x);
  }

  inline C conj() const { return C(x, -y); }
};

const real PI = acosl(-1);
int base = 1;
vector<C> rts = {{0, 0}, {1, 0}};
vector<int> rev = {0, 1};

void ensure_base(int nbase) {
  if (nbase <= base)
    return;
  rev.resize(1 << nbase);
  rts.resize(1 << nbase);
  for (int i = 0; i < (1 << nbase); i++) {
    rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (nbase - 1));
  }
  while (base < nbase) {
    real angle = PI * 2.0 / (1 << (base + 1));
    for (int i = 1 << (base - 1); i < (1 << base); i++) {
      rts[i << 1] = rts[i];
      real angle_i = angle * (2 * i + 1 - (1 << base));
      rts[(i << 1) + 1] = C(cos(angle_i), sin(angle_i));
    }
    ++base;
  }
}

void fft(vector<C> &a, int n) {
  assert((n & (n - 1)) == 0);
  int zeros = __builtin_ctz(n);
  ensure_base(zeros);
  int shift = base - zeros;
  for (int i = 0; i < n; i++) {
    if (i < (rev[i] >> shift)) {
      swap(a[i], a[rev[i] >> shift]);
    }
  }
  for (int k = 1; k < n; k <<= 1) {
    for (int i = 0; i < n; i += 2 * k) {
      for (int j = 0; j < k; j++) {
        C z = a[i + j + k] * rts[j + k];
        a[i + j + k] = a[i + j] - z;
        a[i + j] = a[i + j] + z;
      }
    }
  }
}

vector<int64_t> multiply(const vector<int> &a, const vector<int> &b) {
  int need = (int)a.size() + (int)b.size() - 1;
  int nbase = 1;
  while ((1 << nbase) < need)
    nbase++;
  ensure_base(nbase);
  int sz = 1 << nbase;
  vector<C> fa(sz);
  for (int i = 0; i < sz; i++) {
    int x = (i < (int)a.size() ? a[i] : 0);
    int y = (i < (int)b.size() ? b[i] : 0);
    fa[i] = C(x, y);
  }
  fft(fa, sz);
  C r(0, -0.25 / (sz >> 1)), s(0, 1), t(0.5, 0);
  for (int i = 0; i <= (sz >> 1); i++) {
    int j = (sz - i) & (sz - 1);
    C z = (fa[j] * fa[j] - (fa[i] * fa[i]).conj()) * r;
    fa[j] = (fa[i] * fa[i] - (fa[j] * fa[j]).conj()) * r;
    fa[i] = z;
  }
  for (int i = 0; i < (sz >> 1); i++) {
    C A0 = (fa[i] + fa[i + (sz >> 1)]) * t;
    C A1 = (fa[i] - fa[i + (sz >> 1)]) * t * rts[(sz >> 1) + i];
    fa[i] = A0 + A1 * s;
  }
  fft(fa, sz >> 1);
  vector<int64_t> ret(need);
  for (int i = 0; i < need; i++) {
    ret[i] = llround(i & 1 ? fa[i >> 1].y : fa[i >> 1].x);
  }
  return ret;
}
}; // namespace FastFourierTransform

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int N;
  cin >> N;
  vector<ll> cnt(P);
  rep(i, N) {
    int a;
    cin >> a;
    cnt[a]++;
  }
  vector<ll> beki(P, 1);
  for (int i = 1; i < P; i++)
    beki[i] = beki[i - 1] * 2 % P;
  vector<int> A(P);
  for (int i = 1; i < P; i++) {
    A[i] = cnt[beki[i]];
  }
  auto B = FastFourierTransform::multiply(A, A);
  int L = B.size();
  ll ans = 0;
  rep(i, L) { ans += beki[i % (P - 1)] * B[i]; }
  for (int i = 1; i < P; i++) {
    ans -= beki[i * 2 % (P - 1)] * A[i];
  }
  ans /= 2;
  cout << ans << ln;
}


Error: TLE

Fixed submission:

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
using ll = long long;
using u64 = uint_fast64_t;
using pii = pair<int, int>;
using pll = pair<long long, long long>;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()
constexpr char ln = '\\n';
constexpr long long MOD = 1000000007;
// constexpr long long MOD = 998244353;
template <class T1, class T2> inline bool chmax(T1 &a, T2 b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T1, class T2> inline bool chmin(T1 &a, T2 b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
inline int popcount(int x) { return __builtin_popcount(x); }
inline int popcount(long long x) { return __builtin_popcountll(x); }
void print() { cout << "\\n"; }
template <class T, class... Args> void print(const T &x, const Args &...args) {
  cout << x << " ";
  print(args...);
}
///////////////////////////////////////////////////////////////////////////////////////////////////////////////

const int P = 200003;
const double PI = acos(-1.0);

namespace FastFourierTransform {
using Real = double;
using Complex = complex<Real>;

vector<Complex> DFT(const vector<Complex> &A) {
  int N = A.size();
  if (N == 1)
    return A;

  vector<Complex> even(N / 2), odd(N / 2);

  for (int i = 0; i < N / 2; ++i) {
    even[i] = A[i << 1 | 0];
    odd[i] = A[i << 1 | 1];
  }

  auto inversed_even = DFT(even);
  auto inversed_odd = DFT(odd);

  vector<Complex> ret(N);
  for (int i = 0; i < N; ++i) {
    Complex z = Complex(cos(2 * PI * i / N), sin(2 * PI * i / N));
    ret[i] = inversed_even[i % (N / 2)] + z * inversed_odd[i % (N / 2)];
  }

  return ret;
}

vector<Complex> IDFT(vector<Complex> &A) {
  int N = A.size();

  auto arranged = DFT(A);

  vector<Complex> ret(N);
  for (int i = 0; i < N; ++i)
    ret[i] = arranged[(N - i) % N] / Complex(N, 0);

  return ret;
}

vector<long long> multiply(const vector<int> &A, const vector<int> &B) {
  int sz = (int)A.size() + (int)B.size() - 1;
  int N = 1;
  while (N < sz)
    N *= 2;
  vector<Complex> fA(N), fB(N);
  for (int i = 0; i < A.size(); ++i)
    fA[i] = A[i];
  for (int i = 0; i < B.size(); ++i)
    fB[i] = B[i];

  fA = DFT(fA);
  fB = DFT(fB);

  vector<Complex> C(N);
  for (int i = 0; i < N; ++i)
    C[i] = fA[i] * fB[i];
  C = IDFT(C);
  vector<long long> ret(N);
  for (int i = 0; i < N; ++i)
    ret[i] = (C[i].real() + .1);

  return ret;
}
} // namespace FastFourierTransform

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int N;
  cin >> N;
  vector<ll> cnt(P);
  rep(i, N) {
    int a;
    cin >> a;
    cnt[a]++;
  }
  vector<ll> beki(P, 1);
  for (int i = 1; i < P; i++)
    beki[i] = beki[i - 1] * 2 % P;
  vector<int> A(P);
  for (int i = 1; i < P; i++) {
    A[i] = cnt[beki[i]];
  }
  auto B = FastFourierTransform::multiply(A, A);
  int L = B.size();
  ll ans = 0;
  rep(i, L) { ans += beki[i % (P - 1)] * B[i]; }
  for (int i = 1; i < P; i++) {
    ans -= beki[i * 2 % (P - 1)] * A[i];
  }
  ans /= 2;
  cout << ans << ln;
}