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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <deque>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
#define llint long long
#define inf 1e18
#define rep(x, s, t) for (llint(x) = (s); (x) < (t); (x)++)
#define Rep(x, s, t) for (llint(x) = (s); (x) <= (t); (x)++)
#define chmin(x, y) (x) = min((x), (y))
#define chmax(x, y) (x) = max((x), (y))

using namespace std;
typedef pair<llint, llint> P;

llint n;
llint h = 5, mod = 200003;
vector<llint> vec;
llint a[1 << 19], b[1 << 19];
complex<double> f[1 << 19], F[1 << 19];
llint l[200005];

#define PI 3.141592653589793238462643383

int rev(int x, int n) {
  int ret = 0;
  for (int i = 0; i < n; i++) {
    ret <<= 1;
    ret |= (x >> i) & 1;
  }
  return ret;
}

// f[]とF[]は異なる実体を持たなければならない
void DFT(complex<double> f[], complex<double> F[], int n) {
  int N = 1 << n;
  for (int i = 0; i < N; i++)
    F[rev(i, n)] = f[i];

  complex<double> a, b, x, z;
  for (int i = 1; i <= n; i++) {
    int l = 1 << i;
    z = complex<double>(cos(2 * PI / l), sin(2 * PI / l));
    for (int j = 0; j < N / l; j++) {
      x = complex<double>(1, 0);
      for (int k = 0; k < l / 2; k++) {
        a = F[j * l + k], b = F[j * l + k + l / 2];
        F[j * l + k] = a + x * b;
        F[j * l + k + l / 2] = a - x * b;
        x *= z;
      }
    }
  }
}

// f[]とF[]は異なる実体を持たなければならない
void IDFT(complex<double> F[], complex<double> f[], int n) {
  int N = 1 << n;
  for (int i = 0; i < N; i++)
    f[rev(i, n)] = F[i];

  complex<double> a, b, x, z;
  for (int i = 1; i <= n; i++) {
    int l = 1 << i;
    z = complex<double>(cos(2 * PI / l), -sin(2 * PI / l));
    for (int j = 0; j < N / l; j++) {
      x = complex<double>(1, 0);
      for (int k = 0; k < l / 2; k++) {
        a = f[j * l + k], b = f[j * l + k + l / 2];
        f[j * l + k] = a + x * b;
        f[j * l + k + l / 2] = a - x * b;
        x *= z;
      }
    }
  }
  for (int i = 0; i < N; i++)
    f[i] /= N;
}

llint round(complex<double> c) { return (llint)(c.real() + 0.5); }

llint modpow(llint a, llint n) {
  if (n == 0)
    return 1;
  if (n % 2) {
    return ((a % mod) * (modpow(a, n - 1) % mod)) % mod;
  } else {
    return modpow((a * a) % mod, n / 2) % mod;
  }
}

map<llint, llint> mp;
llint BabyStep_GiantStep(llint a, llint y) {
  mp.clear();
  llint m = sqrt(mod) + 1, mul = 1;
  for (int j = 0; j < m; j++) {
    if (!mp.count(mul))
      mp[mul] = j;
    mul *= a, mul %= mod;
  }
  llint g = modpow(a, m);
  g = modpow(g, mod - 2);

  mul = y;
  for (int i = 0; i < m; i++) {
    if (mp.count(mul))
      return i * m + mp[mul];
    mul *= g, mul %= mod;
  }
  return -1;
}

int main(void) {
  ios::sync_with_stdio(0);
  cin.tie(0);

  cin >> n;
  llint x;
  for (int i = 1; i <= n; i++) {
    cin >> x;
    if (x == 0)
      continue;
    vec.push_back(x);
  }
  n = vec.size();

  llint mul = 1;
  for (int i = 0; i < mod; i++) {
    l[mul] = i;
    mul *= h, mul %= mod;
  }

  for (int i = 0; i < n; i++) {
    llint res = l[vec[i]]; // BabyStep_GiantStep(h, vec[i]);
    a[res]++;
    // cout << vec[i] << " " << res << endl;
  }

  for (int i = 0; i < (1 << 19); i++)
    f[i] = a[i];
  DFT(f, F, 19);
  for (int i = 0; i < (1 << 19); i++)
    F[i] *= F[i];
  IDFT(F, f, 19);
  for (int i = 0; i < (1 << 19); i++)
    a[i] = round(f[i]);

  llint ans = 0;
  for (int i = 0; i < (1 << 19); i++)
    ans += modpow(h, i) * a[i];
  for (int i = 0; i < vec.size(); i++)
    ans -= vec[i] * vec[i] % mod;
  ans /= 2;

  cout << ans << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <deque>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
#define llint long long
#define inf 1e18
#define rep(x, s, t) for (llint(x) = (s); (x) < (t); (x)++)
#define Rep(x, s, t) for (llint(x) = (s); (x) <= (t); (x)++)
#define chmin(x, y) (x) = min((x), (y))
#define chmax(x, y) (x) = max((x), (y))

using namespace std;
typedef pair<llint, llint> P;

llint n;
llint h = 5, mod = 200003;
vector<llint> vec;
llint a[1 << 19], b[1 << 19];
complex<double> f[1 << 19], F[1 << 19];

#define PI 3.141592653589793238

int rev(int x, int n) {
  int ret = 0;
  for (int i = 0; i < n; i++) {
    ret <<= 1;
    ret |= (x >> i) & 1;
  }
  return ret;
}

// f[]とF[]は異なる実体を持たなければならない
void DFT(complex<double> f[], complex<double> F[], int n) {
  int N = 1 << n;
  for (int i = 0; i < N; i++)
    F[rev(i, n)] = f[i];

  complex<double> a, b, x, z;
  for (int i = 1; i <= n; i++) {
    int l = 1 << i;
    z = complex<double>(cos(2 * PI / l), sin(2 * PI / l));
    for (int j = 0; j < N / l; j++) {
      x = complex<double>(1, 0);
      for (int k = 0; k < l / 2; k++) {
        a = F[j * l + k], b = F[j * l + k + l / 2];
        F[j * l + k] = a + x * b;
        F[j * l + k + l / 2] = a - x * b;
        x *= z;
      }
    }
  }
}

// f[]とF[]は異なる実体を持たなければならない
void IDFT(complex<double> F[], complex<double> f[], int n) {
  int N = 1 << n;
  for (int i = 0; i < N; i++)
    f[rev(i, n)] = F[i];

  complex<double> a, b, x, z;
  for (int i = 1; i <= n; i++) {
    int l = 1 << i;
    z = complex<double>(cos(2 * PI / l), -sin(2 * PI / l));
    for (int j = 0; j < N / l; j++) {
      x = complex<double>(1, 0);
      for (int k = 0; k < l / 2; k++) {
        a = f[j * l + k], b = f[j * l + k + l / 2];
        f[j * l + k] = a + x * b;
        f[j * l + k + l / 2] = a - x * b;
        x *= z;
      }
    }
  }
  for (int i = 0; i < N; i++)
    f[i] /= N;
}

llint round(complex<double> c) { return (int)(c.real() + 0.5); }

llint modpow(llint a, llint n) {
  if (n == 0)
    return 1;
  if (n % 2) {
    return ((a % mod) * (modpow(a, n - 1) % mod)) % mod;
  } else {
    return modpow((a * a) % mod, n / 2) % mod;
  }
}

map<llint, llint> mp;
llint BabyStep_GiantStep(llint a, llint y) {
  mp.clear();
  llint m = sqrt(mod) + 1, mul = 1;
  for (int j = 0; j < m; j++) {
    if (!mp.count(mul))
      mp[mul] = j;
    mul *= a, mul %= mod;
  }
  llint g = modpow(a, m);
  g = modpow(g, mod - 2);

  mul = y;
  for (int i = 0; i < m; i++) {
    if (mp.count(mul))
      return i * m + mp[mul];
    mul *= g, mul %= mod;
  }
  return -1;
}

int main(void) {
  ios::sync_with_stdio(0);
  cin.tie(0);

  cin >> n;
  llint x;
  for (int i = 1; i <= n; i++) {
    cin >> x;
    if (x == 0)
      continue;
    vec.push_back(x);
  }
  n = vec.size();

  for (int i = 0; i < n; i++) {
    llint res = BabyStep_GiantStep(h, vec[i]);
    a[res]++;
    // cout << vec[i] << " " << res << endl;
  }

  for (int i = 0; i < (1 << 19); i++)
    f[i] = a[i];
  DFT(f, F, 19);
  for (int i = 0; i < (1 << 19); i++)
    F[i] *= F[i];
  IDFT(F, f, 19);
  for (int i = 0; i < (1 << 19); i++)
    a[i] = round(f[i]);

  llint ans = 0;
  for (int i = 0; i < (1 << 19); i++)
    ans += modpow(h, i) * a[i];
  for (int i = 0; i < vec.size(); i++)
    ans -= vec[i] * vec[i] % mod;
  ans /= 2;

  cout << ans << endl;

  return 0;
}