Problem p02590 (C++) — buggy submission, verdict: Runtime Error

// #include<bits/stdc++.h>
#include <complex>
#include <cstring>
#include <iostream>
#include <unordered_map>
using namespace std;
#define rep(i, a, b) for (int i = (a); i <= (b); ++i)
#define rrep(i, a, b) for (int i = (a); i >= (b); --i)
#define PB push_back
#define ar2 array<int, 2>
typedef long long LL;
typedef complex<double> cp;
const LL P = 200003;
const LL INF = 1e18;
const int N = 524300, Q = sqrt(P);
const double PI = acos(-1);
// mt19937 rng(time(0));

int n, m, x[N], root, p[N], degree[N], id[N];
LL cnt[N];
cp a[N];

LL pw(LL x, LL n) {
  LL res = 1;
  for (; n; n >>= 1, x = x * x % P)
    if (n & 1)
      res = res * x % P;
  return res;
}

void calcDgr() {
  unordered_map<LL, int> pws;
  LL tmp = 1;
  rep(i, 0, Q) pws[tmp] = i, tmp = tmp * root % P;
  rep(x, 2, P - 1) if (p[x] == x) rep(i, 0, P / Q + 1) {
    tmp = pw(root, P - 1 - i * Q % (P - 1)) * x % P;
    if (pws[tmp]) {
      degree[x] = pws[tmp] + i * Q;
      break;
    }
  }
  pws[1] = 0;
  rep(x, 2, P - 1) degree[x] = (degree[x / p[x]] + degree[p[x]]) % (P - 1);
}

void fft(cp a[], int n, int t) {
  rep(i, 0, n - 1) if (i < id[i])
      swap(a[i], a[id[i]]); // 逆变换的时候再swap一次，正好变回原顺序
  for (int h = 2; h <= n; h <<= 1) {
    cp wn(cos(2 * PI / h), t * sin(2 * PI / h));
    for (int j = 0; j < n; j += h) {
      cp w(1, 0), tmp;
      rep(k, j, j + (h >> 1) - 1) {
        tmp = a[k + (h >> 1)] * w;
        a[k + (h >> 1)] = a[k] - tmp;
        a[k] += tmp;
        w *= wn;
      }
    }
  }
  if (t == -1)
    rep(i, 0, n - 1) a[i] /= n;
}

int main() {
  // find the primitive root
  for (LL x = 2;; ++x) {
    bool flag = true;
    for (LL i = 2; i * i <= P - 1; ++i)
      if ((P - 1) % i == 0 && pw(x, (P - 1) / i) == 1) {
        flag = false;
        break;
      }
    if (flag) {
      // cout<<"primitive root = "<<x<<endl;
      root = x;
      break;
    }
  }

  rep(i, 2, P) p[i] = i;
  rep(i, 2, P) if (p[i] == i) for (int j = i + i; j <= P; j += i) p[j] = i;
  calcDgr();
  // rep(x,1,10)cout<<degree[x]<<",";cout<<endl;

  scanf("%d", &n);
  rep(i, 1, n) {
    scanf("%d", &x[i]);
    // cout<<x[i]<<" "<<degree[x[i]]<<endl;
    if (x[i])
      cnt[degree[x[i]]]++;
  }
  for (m = 1; m <= (P << 1); m <<= 1)
    ;
  rep(i, 0, P) a[i] = cnt[i];
  rep(i, 0, m - 1) id[i] = (id[i >> 1] >> 1) | (i & 1 ? m >> 1 : 0);

  fft(a, m, 1);
  rep(i, 0, m - 1) a[i] *= a[i];
  fft(a, m, -1);

  memset(cnt, 0, sizeof cnt);
  rep(i, 0, m) cnt[i % (P - 1)] += (LL)(a[i].real() + .5);
  LL ans = 0;
  rep(i, 0, P - 2) {
    // if(cnt[i]) cout<<i<<" "<<cnt[i]<<endl;
    ans += 1ll * cnt[i] * pw(root, i);
  }
  rep(i, 1, n) ans -= 1ll * x[i] * x[i] % P;

  cout << ans / 2ll << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define rep(i, a, b) for (int i = (a); i <= (b); ++i)
#define rrep(i, a, b) for (int i = (a); i >= (b); --i)
#define PB push_back
#define ar2 array<int, 2>
typedef long long LL;
typedef complex<double> cp;
const LL P = 200003;
const LL INF = 1e18;
const int N = 524300, Q = sqrt(P);
const double PI = acos(-1);
mt19937 rng(time(0));

int n, m, x[N], root, cnt[N], p[N], degree[N], id[N];
cp a[N];

LL pw(LL x, LL n) {
  LL res = 1;
  for (; n; n >>= 1, x = x * x % P)
    if (n & 1)
      res = res * x % P;
  return res;
}

LL calcDgr() {
  unordered_map<LL, int> pws;
  LL tmp = 1;
  rep(i, 0, Q) pws[tmp] = i, tmp = tmp * root % P;
  rep(x, 2, P - 1) if (p[x] == x) rep(i, 0, P / Q + 1) {
    tmp = pw(root, P - 1 - i * Q % (P - 1)) * x % P;
    if (pws[tmp]) {
      degree[x] = pws[tmp] + i * Q;
      break;
    }
  }
  pws[1] = 0;
  rep(x, 2, P - 1) degree[x] = (degree[x / p[x]] + degree[p[x]]) % (P - 1);
}

void fft(cp a[], int n, int t) {
  rep(i, 0, n - 1) if (i < id[i])
      swap(a[i], a[id[i]]); // 逆变换的时候再swap一次，正好变回原顺序
  for (int h = 2; h <= n; h <<= 1) {
    cp wn(cos(2 * PI / h), t * sin(2 * PI / h));
    for (int j = 0; j < n; j += h) {
      cp w(1, 0), tmp;
      rep(k, j, j + (h >> 1) - 1) {
        tmp = a[k + (h >> 1)] * w;
        a[k + (h >> 1)] = a[k] - tmp;
        a[k] += tmp;
        w *= wn;
      }
    }
  }
  if (t == -1)
    rep(i, 0, n - 1) a[i] /= n;
}

int main() {
  // find the primitive root
  for (LL x = 2;; ++x) {
    bool flag = true;
    for (LL i = 2; i * i <= P - 1; ++i)
      if ((P - 1) % i == 0 && pw(x, (P - 1) / i) == 1) {
        flag = false;
        break;
      }
    if (flag) {
      // cout<<"primitive root = "<<x<<endl;
      root = x;
      break;
    }
  }

  rep(i, 2, P) p[i] = i;
  rep(i, 2, P) if (p[i] == i) for (int j = i + i; j <= N; j += i) p[j] = i;
  calcDgr();
  // rep(x,1,10)cout<<degree[x]<<",";cout<<endl;

  scanf("%d", &n);
  rep(i, 1, n) {
    scanf("%d", &x[i]);
    // cout<<x[i]<<" "<<degree[x[i]]<<endl;
    if (x[i])
      cnt[degree[x[i]]]++;
  }
  for (m = 1; m <= (P << 1); m <<= 1)
    ;
  rep(i, 0, P) a[i] = cnt[i];
  rep(i, 0, m - 1) id[i] = (id[i >> 1] >> 1) | (i & 1 ? m >> 1 : 0);

  fft(a, m, 1);
  rep(i, 0, m) a[i] *= a[i];
  fft(a, m, -1);

  memset(cnt, 0, sizeof cnt);
  rep(i, 0, m) cnt[i % (P - 1)] += (LL)(a[i].real() + .5);
  LL ans = 0;
  rep(i, 0, P - 2) {
    // if(cnt[i]) cout<<i<<" "<<cnt[i]<<endl;
    ans += 1ll * cnt[i] * pw(root, i);
  }
  rep(i, 1, n) ans -= 1ll * x[i] * x[i] % P;

  cout << ans / 2 << endl;
  return 0;
}