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

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
// typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;
Int MOD2 = 998244353;
Int P = 200003;
Int order[200004];
Int cnt[200004];
Int cntr[200004];
Int sum, sqsum;
Int n, a;

Int mod_pow(Int x, Int a, Int m = MOD) {
  if (a == 0)
    return 1;
  if (a % 2 == 1)
    return x * mod_pow(x, a - 1, m) % m;
  Int res = mod_pow(x, a / 2, m);
  return res * res % m;
}

Int inv(Int x, Int m = MOD) { return mod_pow(x, m - 2, m); }

void _NTT(vector<Int> &a, vector<Int> &tmp, int l, int r, Int w = -1) {
  if (r - l == 1)
    return;
  if (w == -1) {
    w = mod_pow(3, (MOD2 - 1) / (r - l), MOD2);
  }
  int n = r - l;
  for (int i = 0; i < n / 2; i++) {
    tmp[l + i] = a[l + i * 2];
    tmp[l + n / 2 + i] = a[l + i * 2 + 1];
  }
  _NTT(tmp, a, l, l + n / 2, w * w % MOD2);
  _NTT(tmp, a, l + n / 2, r, w * w % MOD2);

  Int wi = 1;
  for (int i = 0; i < n; i++) {
    a[l + i] = tmp[l + i % (n / 2)] + tmp[l + n / 2 + i % (n / 2)] * wi;
    a[l + i] %= MOD2;
    wi = wi * w % MOD2;
  }
}

vector<Int> NTT(vector<Int> &a) {
  vector<Int> res = a;
  vector<Int> tmp = a;
  _NTT(res, tmp, 0, a.size());
  return res;
}

vector<Int> InvNTT(vector<Int> &a) {
  reverse(next(a.begin()), a.end());
  vector<Int> res = NTT(a);
  reverse(next(a.begin()), a.end());
  Int invn = inv(a.size(), MOD2);
  for (int i = 0; i < res.size(); i++)
    res[i] = res[i] * invn % MOD2;
  return res;
}

vector<Int> conv(vector<Int> &a, vector<Int> &b) {
  int n = a.size();
  auto A = NTT(a);
  auto B = NTT(b);
  for (int i = 0; i < n; i++)
    A[i] = A[i] * B[i] % MOD2;
  auto res = InvNTT(A);
  return res;
}

int main() {
  Int now = 1;
  bool ok = true;
  Int extra = 0;
  for (Int i = 0; i < P - 1; i++) {
    order[now] = i;
    now = now * 2 % P;
  }
  vector<Int> cnt(1 << 19, 0);
  cin >> n;
  for (int i = 0; i < n; i++) {
    cin >> a;
    extra += a * a % P;
    if (a == 0)
      continue;
    cnt[order[a]] += 1;
  }
  auto CNT = NTT(cnt);
  for (auto &elem : CNT)
    elem = elem * elem % MOD2;
  auto result1 = InvNTT(CNT);
  Int mod2 = MOD2;
  MOD2 = 469762049;
  auto CNT2 = NTT(cnt);
  for (auto &elem : CNT2)
    elem = elem * elem % MOD2;
  auto result2 = InvNTT(CNT2);

  Int ans = 0;
  for (Int i = 0, now = 1; i < (1 << 19); i++) {
    Int r1 = result1[i], r2 = result2[i];
    if (r1)
      while (r1 % MOD2 != r2)
        r1 += mod2;
    ans += now * r1;
    now = now * 2 % P;
  }

  ans -= extra;
  cout << ans / 2 << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define MOD 1000000007
#define MOD2 998244353
#define INF (1 << 29)
#define LINF (1LL << 60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
// typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int P = 200003;
Int order[200004];
Int cnt[200004];
Int cntr[200004];
Int sum, sqsum;
Int n, a;

vector<CP> INV_FFT(vector<CP> &);
vector<CP> FFT2POW(vector<CP> &x, CP w) {
  if (x.size() == 1)
    return x;
  Int n = x.size();
  vector<CP> even;
  vector<CP> odd;
  for (Int i = 0; i < n; i++) {
    if (i % 2 == 0)
      even.push_back(x[i]);
    else
      odd.push_back(x[i]);
  }
  even = FFT2POW(even, w * w);
  odd = FFT2POW(odd, w * w);
  vector<CP> res(n);
  CP w_i = 1;
  for (Int i = 0; i < n; i++) {
    res[i] = even[i % even.size()] + w_i * odd[i % odd.size()];
    w_i *= w;
  }
  return res;
}

CP primitive_root(Int n) {
  double theta = 2 * PI / n;
  return CP(cos(theta), sin(theta));
}

vector<CP> FFT(vector<CP> &a) {
  Int n = a.size();
  if ((n & -n) == n)
    return FFT2POW(a, primitive_root(n));

  Int m = 1;
  while (m < 2 * n)
    m <<= 1;

  CP w = primitive_root(2 * n);
  vector<CP> ws(2 * n + 1, 1);
  for (Int i = 1; i <= 2 * n; i++)
    ws[i] = ws[i - 1] * w;

  vector<CP> winv(m, 0);
  vector<CP> aw(m, 0);
  for (Int i = 0; i < n; i++) {
    Int p = -(i * i) % (2 * n) + 2 * n;
    winv[n - i] = winv[n + i] = ws[p];
    aw[i] = a[i] * ws[i * i % (2 * n)];
  }
  auto AW = FFT(aw);
  auto WINV = FFT(winv);
  for (int i = 0; i < m; i++)
    WINV[i] *= AW[i];
  vector<CP> c = INV_FFT(WINV);
  vector<CP> A(n);
  for (Int i = 0; i < n; i++) {
    A[i] = ws[i * i % (2 * n)] * c[n + i];
  }
  return A;
}

vector<CP> INV_FFT(vector<CP> &x) {
  reverse(x.begin() + 1, x.end());
  auto res = FFT(x);
  for (auto &elem : res)
    elem /= res.size();
  return res;
}

int main() {
  Int now = 1;
  bool ok = true;
  Int extra = 0;
  for (Int i = 0; i < P - 1; i++) {
    order[now] = i;
    now = now * 2 % P;
  }
  vector<CP> cnt(200002, 0);
  cin >> n;
  for (int i = 0; i < n; i++) {
    cin >> a;
    extra += a * a % P;
    if (a == 0)
      continue;
    cnt[order[a]] += 1;
  }
  auto CNT = FFT(cnt);
  for (auto &elem : CNT)
    elem *= elem;
  auto result = INV_FFT(CNT);
  Int ans = 0;
  now = 1;
  for (Int i = 0; i < P - 1; i++) {
    Int r = Int(result[i].real() + 0.5);
    ans += now * r;
    now = now * 2 % P;
  }
  ans -= extra;
  cout << ans / 2 << endl;
  return 0;
}