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

#include <bits/stdc++.h>

using namespace std;
using int64 = long long;

constexpr int DEBUG = 0;

constexpr int64 P = 200003;

using Complex = complex<double>;

// Reference: http://www.prefield.com/algorithm/math/fft.html
vector<Complex> DFT(const vector<Complex> &p, int n, int sign = 1) {
  const double theta = 2.0 * M_PI / n * sign;
  int unit = 1;
  vector<Complex> q = p;
  for (int m = n; m >= 2; m >>= 1) {
    int mh = m >> 1;
    for (int i = 0; i < mh; i++) {
      Complex w = Complex(cos(theta * unit * i), sin(theta * unit * i));
      for (int j = i; j < n; j += m) {
        int k = j + mh;
        Complex x = q[j] - q[k];
        q[j] += q[k];
        q[k] = w * x;
      }
    }
    unit *= 2;
  }
  int i = 0;
  for (int j = 1; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1)
      ;
    if (j < i)
      swap(q[i], q[j]);
  }
  return q;
}

vector<Complex> InvDFT(const vector<Complex> &p, int n) {
  auto q = DFT(p, n, -1);
  for (int i = 0; i < n; i++) {
    q[i] /= n;
  }
  return q;
}

vector<Complex> MultiplyFFT(const vector<Complex> &p1,
                            const vector<Complex> &p2) {
  int d = p1.size() + p2.size() - 1;
  int n = 1;
  while (n < d + 1)
    n *= 2;

  auto q1 = p1;
  q1.resize(n);
  auto q2 = p2;
  q2.resize(n);
  auto r1 = DFT(q1, n);
  auto r2 = DFT(q2, n);

  vector<Complex> s(n);
  for (int i = 0; i < n; i++) {
    s[i] = r1[i] * r2[i];
  }
  return InvDFT(s, n);
}

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

  // i = 2^p
  vector<int64> i_to_p(P, -1);
  vector<int64> p_to_i(P - 1, -1);
  {
    int64 x = 1;
    for (int i = 0; i < P - 1; i++) {
      i_to_p[x] = i;
      p_to_i[i] = x;
      x *= 2;
      x %= P;
    }
  }

  int n;
  cin >> n;

  int64 delta = 0;
  vector<Complex> xs(P - 1);
  for (int i = 0; i < n; i++) {
    int64 a;
    cin >> a;
    if (a == 0)
      continue;
    xs[i_to_p[a]] += 1;
    // cout << a * a % P << endl;
    delta += a * a % P;
  }
  // cout << "Before FFT" << endl;
  vector<Complex> ys = MultiplyFFT(xs, xs);
  // cout << "After FFT" << endl;
  int64 ans = 0;
  for (int i = 0; i < ys.size(); i++) {
    int k = i % (P - 1);
    // int64 c = roundf64(ys[i].real());
    // int64 v = p_to_i[k];
    // if (c > 0) {
    //   cout << c << " " << v << endl;
    // }
    ans += p_to_i[k] * roundf64(ys[i].real());
  }
  // cout << ans << endl;
  ans -= delta;
  ans /= 2;
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using int64 = long long;

constexpr int DEBUG = 0;

constexpr int64 P = 200003;

using Complex = complex<double>;

vector<Complex> DFT(const vector<Complex> &p, int n, int sign = 1) {
  if (n == 1)
    return p;

  vector<Complex> q1(n / 2), q2(n / 2);
  for (int i = 0; i < n / 2; i++) {
    q1[i] = p[2 * i + 0];
    q2[i] = p[2 * i + 1];
  }

  auto r1 = DFT(q1, n / 2, sign);
  auto r2 = DFT(q2, n / 2, sign);

  vector<Complex> s(n);
  for (int i = 0; i < n; i++) {
    double theta = 2.0 * M_PI * i / n * sign;
    s[i] = r1[i % (n / 2)] + Complex(cos(theta), sin(theta)) * r2[i % (n / 2)];
  }
  return s;
}

vector<Complex> InvDFT(const vector<Complex> &p, int n) {
  auto q = DFT(p, n, -1);
  for (int i = 0; i < n; i++) {
    q[i] /= n;
  }
  return q;
}

// Verified: ATC001-C
vector<Complex> MultiplyFFT(const vector<Complex> &p1,
                            const vector<Complex> &p2) {
  int d = p1.size() + p2.size() - 1;
  int n = 1;
  while (n < d + 1)
    n *= 2;

  auto q1 = p1;
  q1.resize(n);
  auto q2 = p2;
  q2.resize(n);
  auto r1 = DFT(q1, n);
  auto r2 = DFT(q2, n);

  vector<Complex> s(n);
  for (int i = 0; i < n; i++) {
    s[i] = r1[i] * r2[i];
  }
  return InvDFT(s, n);
}

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

  // i = 2^p
  vector<int64> i_to_p(P, -1);
  vector<int64> p_to_i(P - 1, -1);
  {
    int64 x = 1;
    for (int i = 0; i < P - 1; i++) {
      i_to_p[x] = i;
      p_to_i[i] = x;
      x *= 2;
      x %= P;
    }
  }

  int n;
  cin >> n;

  int64 delta = 0;
  vector<Complex> xs(P - 1);
  for (int i = 0; i < n; i++) {
    int64 a;
    cin >> a;
    if (a == 0)
      continue;
    xs[i_to_p[a]] += 1;
    // cout << a * a % P << endl;
    delta += a * a % P;
  }
  // cout << "Before FFT" << endl;
  vector<Complex> ys = MultiplyFFT(xs, xs);
  // cout << "After FFT" << endl;
  int64 ans = 0;
  for (int i = 0; i < ys.size(); i++) {
    int k = i % (P - 1);
    // int64 c = roundf64(ys[i].real());
    // int64 v = p_to_i[k];
    // if (c > 0) {
    //   cout << c << " " << v << endl;
    // }
    ans += p_to_i[k] * roundf64(ys[i].real());
  }
  // cout << ans << endl;
  ans -= delta;
  ans /= 2;
  cout << ans << endl;
}