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

#include <bits/stdc++.h>

typedef long long int64;
typedef std::complex<double> C;

const int N = 4e5, MOD = 2e5 + 3, G = 2;
const double PI = acos(-1);

int n;
int64 cnt[N + 5];
std::unordered_map<int, int> map;

inline int power(int x, int k, int mod) {
  int ans = 1;
  for (; k > 0; k >>= 1, x = (int64)x * x % mod) {
    if (k & 1) {
      ans = (int64)ans * x % mod;
    }
  }
  return ans;
}
inline int inv(int x, int mod) { return pow(x, mod - 2); }
void fft(C *a, int k, bool opt) {
  int n = 1 << k;
  std::vector<int> rev(n);
  for (int i = 0; i < n; i++) {
    rev[i] = rev[i >> 1] >> 1 | (i & 1) << (k - 1);
    if (i < rev[i]) {
      std::swap(a[i], a[rev[i]]);
    }
  }
  for (int l = 1; l < n; l <<= 1) {
    C w = C(cos(PI / l), sin(PI / l));
    for (int i = 0; i < n; i += l << 1) {
      C wk = C(1, 0);
      for (int j = 0; j < l; j++, wk *= w) {
        C x = a[i + j + l] * wk;
        a[i + j + l] = a[i + j] - x;
        a[i + j] += x;
      }
    }
  }
  if (opt) {
    std::reverse(a + 1, a + n);
    for (int i = 0; i < n; i++) {
      a[i] /= n;
    }
  }
}
void dft(C *a, int k) { fft(a, k, false); }
void idft(C *a, int k) { fft(a, k, true); }
void selfMul(int64 *a) {
  static C b[1 << 19];
  int k = 19, n = 1 << k;
  std::fill(b, b + n, C(0, 0));
  for (int i = 0; i < MOD; i++) {
    b[i] = C(a[i], 0);
  }
  dft(b, k);
  for (int i = 0; i < n; i++) {
    b[i] *= b[i];
  }
  idft(b, k);
  for (int i = 0; i < MOD * 2; i++) {
    a[i] = (int64)(b[i].real() + 0.5);
  }
}
int main() {
  scanf("%d", &n);
  for (int w = 1, i = 0; i < MOD; i++, w = (int64)w * G % MOD) {
    map[w] = i;
  }
  int64 ans = 0;
  for (int x, i = 1; i <= n; i++) {
    scanf("%d", &x);
    if (x > 0) {
      ans -= (int64)x * x % MOD;
      cnt[map[x]]++;
    }
  }
  selfMul(cnt);
  for (int w = 1, i = 0; i < MOD * 2; i++, w = (int64)w * G % MOD) {
    ans += cnt[i] * w;
  }
  printf("%lld\\n", ans / 2);
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

typedef long long int64;
typedef std::complex<double> C;

const int N = 4e5, MOD = 2e5 + 3, G = 2;
const double PI = acos(-1);

int n;
int64 cnt[N + 5];

inline int power(int x, int k, int mod) {
  int ans = 1;
  for (; k > 0; k >>= 1, x = (int64)x * x % mod) {
    if (k & 1) {
      ans = (int64)ans * x % mod;
    }
  }
  return ans;
}
inline int inv(int x, int mod) { return pow(x, mod - 2); }
void fft(C *a, int k, bool opt) {
  int n = 1 << k;
  std::vector<int> rev(n);
  for (int i = 0; i < n; i++) {
    rev[i] = rev[i >> 1] >> 1 | (i & 1) << (k - 1);
    if (i < rev[i]) {
      std::swap(a[i], a[rev[i]]);
    }
  }
  for (int l = 1; l < n; l <<= 1) {
    C w = C(cos(PI / l), sin(PI / l));
    for (int i = 0; i < n; i += l << 1) {
      C wk = C(1, 0);
      for (int j = 0; j < l; j++, wk *= w) {
        C x = a[i + j + l] * wk;
        a[i + j + l] = a[i + j] - x;
        a[i + j] += x;
      }
    }
  }
  if (opt) {
    std::reverse(a + 1, a + n);
    for (int i = 0; i < n; i++) {
      a[i] /= n;
    }
  }
}
void dft(C *a, int k) { fft(a, k, false); }
void idft(C *a, int k) { fft(a, k, true); }
void selfMul(int64 *a) {
  static C b[1 << 19];
  int k = 19, n = 1 << k;
  std::fill(b, b + n, C(0, 0));
  for (int i = 0; i < MOD; i++) {
    b[i] = C(a[i], 0);
  }
  dft(b, k);
  for (int i = 0; i < n; i++) {
    b[i] *= b[i];
  }
  idft(b, k);
  for (int i = 0; i < MOD * 2; i++) {
    a[i] = (int64)(b[i].real() + 0.5);
  }
}
inline int calc(int a, int b) {
  if (b == 1) {
    return 0;
  }
  std::unordered_map<int, int> map;
  // 2 ^ ??? = x (mod p)
  int m = ceil(sqrt(MOD)), w = 1;
  for (int i = 0; i < m; i++, b = (int64)b * a % MOD, w = (int64)w * a % MOD) {
    map[b] = i;
  }
  for (int v = w, i = 1; i <= m; i++, v = (int64)v * w % MOD) {
    if (map.count(v)) {
      return i * m - map[v];
    }
  }
  assert(false);
}
int main() {
  scanf("%d", &n);
  int64 ans = 0;
  for (int x, i = 1; i <= n; i++) {
    scanf("%d", &x);
    if (x > 0) {
      ans -= (int64)x * x % MOD;
      cnt[calc(G, x)]++;
    }
  }
  selfMul(cnt);
  for (int w = 1, i = 0; i < MOD * 2; i++, w = (int64)w * G % MOD) {
    ans += cnt[i] * w;
  }
  printf("%lld\\n", ans / 2);
  return 0;
}