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

// chrono::system_clock::now().time_since_epoch().count()
#include <bits/stdc++.h>

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define debug(x) cerr << #x << " = " << x << endl;

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef complex<double> cd;

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

int addMod(int a, int b, int m) {
  a += b;

  if (m <= a) {
    a -= m;
  }

  return a;
}

int mulMod(int a, int b, int m) { return a * 1ll * b % m; }

int binPow(int a, int b, int m) {
  int ret = 1;

  while (b > 0) {
    if (b & 1) {
      ret = mulMod(ret, a, m);
    }

    a = mulMod(a, a, m);
    b >>= 1;
  }

  return ret;
}

int n, arr[P], mem[P];
ll cnt[P], cnt2[2 * P];
cd tmp[3 * P];
int pw[2 * P];

ll f(ll x) { return x * (x - 1) / 2; }

void fft(cd p[], int n, bool invert) {
  int k = 0;

  while ((1 << k) < n) {
    ++k;
  }

  assert((1 << k) == n);

  for (int i = 0; i < n; ++i) {
    int i2 = 0;

    for (int j = 0; j < k; ++j) {
      if ((i >> j) & 1) {
        i2 += (1 << (k - j - 1));
      }
    }

    if (i2 < i) {
      swap(p[i], p[i2]);
    }
  }

  for (int len = 2; len <= n; len <<= 1) {
    double ang = 2 * PI / len * (invert ? -1 : 1);
    cd wlen(cos(ang), sin(ang));

    for (int i = 0; i < n; i += len) {
      cd w(1);

      for (int j = 0; j < len / 2; ++j) {
        cd u = p[i + j], v = p[i + j + len / 2] * w;
        p[i + j] = u + v;
        p[i + j + len / 2] = u - v;
        w *= wlen;
      }
    }
  }

  if (invert) {
    for (int i = 0; i < n; ++i) {
      p[i] /= n;
    }
  }
}

void solve() {
  scanf("%d", &n);

  for (int i = 1, j = 0; i <= n; ++i) {
    int x;
    scanf("%d", &x);

    if (x) {
      arr[++j] = x;
    }
  }

  while (n > 0 && arr[n] == 0) {
    --n;
  }

  fill(mem, mem + P, -1);

  for (int i = 0, cur = 1; i < 2 * P; ++i) {
    if (mem[cur] == -1) {
      mem[cur] = i;
    }

    pw[i] = cur;
    cur = mulMod(cur, G, P);
  }

  for (int i = 1; i <= n; ++i) {
    ++cnt[mem[arr[i]]];
  }

  for (int i = 0; i < P; ++i) {
    tmp[i] = cd(cnt[i]);
  }

  int m = P;

  while (m & (m - 1)) {
    ++m;
  }

  m += m;
  fft(tmp, m, false);

  for (int i = 0; i < m; ++i) {
    tmp[i] *= tmp[i];
  }

  fft(tmp, m, true);

  for (int i = 0; i < 2 * P; ++i) {
    cnt2[i] = (ll)(tmp[i].real() + 0.5);
  }

  ll ans = 0;

  for (int i = 0; i < 2 * P; ++i) {
    ll cur = cnt2[i];

    if (i % 2 == 0) {
      cur -= cnt[i / 2] * cnt[i / 2];
      cur /= 2;
      cur += f(cnt[i / 2]);
    } else {
      cur /= 2;
    }

    ans += cur * 1ll * pw[i];
  }

  printf("%lld\\n", ans);
}

int main() {
  int tt = 1;

  while (tt--) {
    solve();
  }

  return 0;
}


Error: TLE

Fixed submission:

// chrono::system_clock::now().time_since_epoch().count()
#include <bits/stdc++.h>

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define fi first
#define se second
#define all(x) (x).begin(), (x).end()
#define debug(x) cerr << #x << " = " << x << endl;

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;

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

int addMod(int a, int b, int m) {
  a += b;

  if (m <= a) {
    a -= m;
  }

  return a;
}

int mulMod(int a, int b, int m) { return a * 1ll * b % m; }

int binPow(int a, int b, int m) {
  int ret = 1;

  while (b > 0) {
    if (b & 1) {
      ret = mulMod(ret, a, m);
    }

    a = mulMod(a, a, m);
    b >>= 1;
  }

  return ret;
}

int n, arr[P], mem[P];
ll cnt[P], cnt2[2 * P];
int pw[2 * P];

ll f(ll x) { return x * (x - 1) / 2; }

void solve() {
  scanf("%d", &n);

  for (int i = 1, j = 0; i <= n; ++i) {
    int x;
    scanf("%d", &x);

    if (x) {
      arr[++j] = x;
    }
  }

  while (n > 0 && arr[n] == 0) {
    --n;
  }

  fill(mem, mem + P, -1);

  for (int i = 0, cur = 1; i < 2 * P; ++i) {
    if (mem[cur] == -1) {
      mem[cur] = i;
    }

    pw[i] = cur;
    cur = mulMod(cur, G, P);
  }

  for (int i = 1; i <= n; ++i) {
    ++cnt[mem[arr[i]]];
  }

  vector<int> vec;

  for (int i = 0; i < P; ++i) {
    if (cnt[i] > 0) {
      vec.pb(i);
    }
  }

  for (int a : vec) {
    for (int b : vec) {
      cnt2[a + b] += cnt[a] * cnt[b];
    }
  }

  ll ans = 0;

  for (int i = 0; i < 2 * P; ++i) {
    ll cur = cnt2[i];

    if (i % 2 == 0) {
      cur -= cnt[i / 2] * cnt[i / 2];
      cur /= 2;
      cur += f(cnt[i / 2]);
    } else {
      cur /= 2;
    }

    ans += cur * 1ll * pw[i];
  }

  printf("%lld\\n", ans);
}

int main() {
  int tt = 1;

  while (tt--) {
    solve();
  }

  return 0;
}