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

#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const int INF = 1e9;
const ll LINF = 1e18;
inline ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
inline ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }
template <class S, class T>
ostream &operator<<(ostream &out, const pair<S, T> &o) {
  out << "(" << o.first << "," << o.second << ")";
  return out;
}
template <class T> ostream &operator<<(ostream &out, const vector<T> &V) {
  for (int i = 0; i < V.size(); i++) {
    out << V[i];
    if (i != V.size() - 1)
      out << " ";
  }
  return out;
}
template <class T>
ostream &operator<<(ostream &out, const vector<vector<T>> &Mat) {
  for (int i = 0; i < Mat.size(); i++) {
    if (i != 0)
      out << endl;
    out << Mat[i];
  }
  return out;
}
template <class S, class T>
ostream &operator<<(ostream &out, const map<S, T> &mp) {
  out << "{ ";
  for (auto it = mp.begin(); it != mp.end(); it++) {
    out << it->first << ":" << it->second;
    if (mp.size() - 1 != distance(mp.begin(), it))
      out << ", ";
  }
  out << " }";
  return out;
}
template <typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}
template <typename T, typename V>
typename enable_if<is_class<T>::value == 0>::type fill_v(T &t, const V &v) {
  t = v;
}
template <typename T, typename V>
typename enable_if<is_class<T>::value != 0>::type fill_v(T &t, const V &v) {
  for (auto &e : t)
    fill_v(e, v);
}
/*
 <url:https://atcoder.jp/contests/agc047/tasks/agc047_c>
 問題文============================================================
 C - Product Modulo
 =================================================================
 解説=============================================================
 ================================================================
 */
/*
 Example of use
 // 列a, bを用意
 std::vector<int> a, b;

 // a，bを畳み込んだ結果をcとする
 std::vector<int> c(FFT::mul(a, b));
 */

// Description: 列a[i],b[i]から列c[k]=sum{a[i]*b[k-i]}を生成する.O(NlogN).
namespace FFT {
typedef long long TYPE;
const double pi = std::acos(-1);
std::vector<std::complex<double>> tmp;
size_t sz = 1;

std::vector<std::complex<double>> fft(std::vector<std::complex<double>> a,
                                      bool inv = false) {
  size_t mask = sz - 1, p = 0;
  for (size_t i = sz >> 1; i >= 1; i >>= 1) {
    auto &cur = (p & 1) ? tmp : a;
    auto &nex = (p & 1) ? a : tmp;
    std::complex<double> e = std::polar(
        1., 2 * pi * i * (inv ? -1 : 1) / sz); // (絶対値,偏角)の極形式指定
    std::complex<double> w = 1;
    for (size_t j = 0; j < sz; j += i) {
      for (size_t k = 0; k < i; ++k) {
        nex[j + k] =
            cur[((j << 1) & mask) + k] + w * cur[(((j << 1) + i) & mask) + k];
      }
      w *= e;
    }
    ++p;
  }
  if (p & 1)
    std::swap(a, tmp);
  if (inv)
    for (size_t i = 0; i < sz; ++i)
      a[i] /= sz;
  return a;
}
std::vector<TYPE> mul(std::vector<TYPE> a, std::vector<TYPE> b) {
  size_t m = a.size() + b.size() - 1;
  sz = 1;
  while (m > sz)
    sz <<= 1;
  tmp.clear();
  tmp.resize(sz);
  std::vector<std::complex<double>> A(sz), B(sz);
  for (size_t i = 0; i < a.size(); ++i)
    A[i].real(a[i]);
  for (size_t i = 0; i < b.size(); ++i)
    B[i].real(b[i]);
  A = fft(A);
  B = fft(B);
  for (size_t i = 0; i < sz; ++i)
    A[i] *= B[i];
  A = fft(A, true);
  a.clear();
  a.resize(m);
  for (size_t i = 0; i < m; ++i)
    a[i] = std::round(A[i].real());
  return a;
}
}; // namespace FFT

const ll P = 200'003;
template <class Type> Type solve(Type res = Type()) {
  // primitive root g = 2;
  ll g = 2;
  // vector<ll> order{1};
  vector<ll> order(2 * P);
  order[0] = 1;
  for (int i = 1; i < 2 * P; i++) {
    order[i] = order[i - 1] * g % P;
  }
  // for (ll x = g; x != 1; x = x * g % P) {
  //     order.emplace_back(x);
  // }
  // assert(order.size() == P - 1);

  vector<ll> where(P);
  for (int i = 0; i < P; i++) {
    where[order[i]] = i;
  }

  vector<ll> cnt(P);

  int N;
  cin >> N;
  vector<ll> A(N);
  for (auto &in : A)
    cin >> in;
  for (auto a : A) {
    if (a == 0)
      continue;
    cnt[where[a]]++;
  }

  vector<ll> multiplier(P);
  auto answer = FFT::mul(cnt, cnt);

  for (int i = 0; i < (int)answer.size(); i++) {
    res += order[i] * answer[i];
  }

  for (ll a : A) {
    res -= a * a % P;
  }

  res /= 2;
  return res;
}
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // solve<ll>(0);
  cout << fixed << setprecision(12) << solve<ll>() << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
const int INF = 1e9;
const ll LINF = 1e18;
inline ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
inline ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }
template <class S, class T>
ostream &operator<<(ostream &out, const pair<S, T> &o) {
  out << "(" << o.first << "," << o.second << ")";
  return out;
}
template <class T> ostream &operator<<(ostream &out, const vector<T> &V) {
  for (int i = 0; i < V.size(); i++) {
    out << V[i];
    if (i != V.size() - 1)
      out << " ";
  }
  return out;
}
template <class T>
ostream &operator<<(ostream &out, const vector<vector<T>> &Mat) {
  for (int i = 0; i < Mat.size(); i++) {
    if (i != 0)
      out << endl;
    out << Mat[i];
  }
  return out;
}
template <class S, class T>
ostream &operator<<(ostream &out, const map<S, T> &mp) {
  out << "{ ";
  for (auto it = mp.begin(); it != mp.end(); it++) {
    out << it->first << ":" << it->second;
    if (mp.size() - 1 != distance(mp.begin(), it))
      out << ", ";
  }
  out << " }";
  return out;
}
template <typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}
template <typename T, typename V>
typename enable_if<is_class<T>::value == 0>::type fill_v(T &t, const V &v) {
  t = v;
}
template <typename T, typename V>
typename enable_if<is_class<T>::value != 0>::type fill_v(T &t, const V &v) {
  for (auto &e : t)
    fill_v(e, v);
}
/*
 <url:https://atcoder.jp/contests/agc047/tasks/agc047_c>
 問題文============================================================
 C - Product Modulo
 =================================================================
 解説=============================================================
 ================================================================
 */
/*
 Example of use
 // 列a, bを用意
 std::vector<int> a, b;

 // a，bを畳み込んだ結果をcとする
 std::vector<int> c(FFT::mul(a, b));
 */

// Description: 列a[i],b[i]から列c[k]=sum{a[i]*b[k-i]}を生成する.O(NlogN).
namespace FFT {
typedef long long TYPE;
const double pi = std::acos(-1);
std::vector<std::complex<double>> tmp;
size_t sz = 1;

std::vector<std::complex<double>> fft(std::vector<std::complex<double>> a,
                                      bool inv = false) {
  size_t mask = sz - 1, p = 0;
  for (size_t i = sz >> 1; i >= 1; i >>= 1) {
    auto &cur = (p & 1) ? tmp : a;
    auto &nex = (p & 1) ? a : tmp;
    std::complex<double> e = std::polar(
        1., 2 * pi * i * (inv ? -1 : 1) / sz); // (絶対値,偏角)の極形式指定
    std::complex<double> w = 1;
    for (size_t j = 0; j < sz; j += i) {
      for (size_t k = 0; k < i; ++k) {
        nex[j + k] =
            cur[((j << 1) & mask) + k] + w * cur[(((j << 1) + i) & mask) + k];
      }
      w *= e;
    }
    ++p;
  }
  if (p & 1)
    std::swap(a, tmp);
  if (inv)
    for (size_t i = 0; i < sz; ++i)
      a[i] /= sz;
  return a;
}
std::vector<TYPE> mul(std::vector<TYPE> a, std::vector<TYPE> b) {
  size_t m = a.size() + b.size() - 1;
  sz = 1;
  while (m > sz)
    sz <<= 1;
  tmp.clear();
  tmp.resize(sz);
  std::vector<std::complex<double>> A(sz), B(sz);
  for (size_t i = 0; i < a.size(); ++i)
    A[i].real(a[i]);
  for (size_t i = 0; i < b.size(); ++i)
    B[i].real(b[i]);
  A = fft(A);
  B = fft(B);
  for (size_t i = 0; i < sz; ++i)
    A[i] *= B[i];
  A = fft(A, true);
  a.clear();
  a.resize(m);
  for (size_t i = 0; i < m; ++i)
    a[i] = std::round(A[i].real());
  return a;
}
}; // namespace FFT

const ll P = 200'003;
template <class Type> Type solve(Type res = Type()) {
  // primitive root g = 2;
  ll g = 2;
  vector<ll> order{1};
  for (ll x = g; x != 1; x = x * g % P) {
    order.emplace_back(x);
  }
  assert(order.size() == P - 1);
  vector<ll> where(P);
  for (int i = 0; i < (int)order.size(); i++) {
    where[order[i]] = i;
  }

  vector<ll> cnt(P);

  int N;
  cin >> N;
  vector<ll> A(N);
  for (auto &in : A)
    cin >> in;
  for (auto a : A) {
    if (a == 0)
      continue;
    cnt[where[a]]++;
  }

  vector<ll> multiplier(P);
  auto answer = FFT::mul(cnt, cnt);

  for (int i = 0; i < (int)answer.size(); i++) {
    if (answer[i]) {
      ll product = order[i % order.size()];
      multiplier[product] += answer[i];
    }
  }

  ll sum = 0;
  for (ll a : A) {
    if (a == 0)
      continue;
    --multiplier[a * a % P];
    sum += a * a % P;
  }

  for (int i = 1; i < P; i++) {
    assert(multiplier[i] % 2 == 0);
    res += multiplier[i] / 2 * i;
  }
  return res;
}
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(false);
  // solve<ll>(0);
  cout << fixed << setprecision(12) << solve<ll>() << endl;
  return 0;
}