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

#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);
}

vector<int64> MultiplyFFT(const vector<int64> &p1, const vector<int64> &p2) {
  vector<Complex> q1;
  for (int i = 0; i < p1.size(); i++) {
    q1.emplace_back(p1[i], 0);
  }
  vector<Complex> q2;
  for (int i = 0; i < p2.size(); i++) {
    q2.emplace_back(p2[i], 0);
  }
  vector<Complex> r = MultiplyFFT(q1, q2);
  vector<int64> s;
  for (int i = 0; i < r.size(); i++) {
    s.push_back(llround(r[i].real()));
  }
  return s;
}

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<int64> xs(P - 1);
  for (int i = 0; i < n; i++) {
    int64 a;
    cin >> a;
    if (a == 0)
      continue;
    xs[i_to_p[a]] += 1;
    delta += a * a % P;
  }
  vector<int64> ys = MultiplyFFT(xs, xs);
  int64 ans = 0;
  for (int i = 0; i < ys.size(); i++) {
    int k = i % (P - 1);
    ans += p_to_i[k] * ys[i];
  }
  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>;

// 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);
}

vector<int64> MultiplyFFT(const vector<int64> &p1, const vector<int64> &p2) {
  vector<Complex> q1;
  for (int i = 0; i < p1.size(); i++) {
    q1.emplace_back(p1[i], 0);
  }
  vector<Complex> q2;
  for (int i = 0; i < p2.size(); i++) {
    q2.emplace_back(p2[i], 0);
  }
  vector<Complex> r = MultiplyFFT(q1, q2);
  vector<int64> s;
  for (int i = 0; i < r.size(); i++) {
    s.push_back(llround(r[i].real()));
  }
}

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<int64> xs(P - 1);
  for (int i = 0; i < n; i++) {
    int64 a;
    cin >> a;
    if (a == 0)
      continue;
    xs[i_to_p[a]] += 1;
    delta += a * a % P;
  }
  vector<int64> ys = MultiplyFFT(xs, xs);
  int64 ans = 0;
  for (int i = 0; i < ys.size(); i++) {
    int k = i % (P - 1);
    ans += p_to_i[k] * ys[i];
  }
  ans -= delta;
  ans /= 2;
  cout << ans << endl;
}