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


// Problem : C - Sum of product of pairs
// Contest : AtCoder - AtCoder Beginner Contest 177
// URL : https://atcoder.jp/contests/abc177/tasks/abc177_c
// Memory Limit : 1024 MB
// Time Limit : 2000 ms
// Powered by CP Editor (https://github.com/cpeditor/cpeditor)

// 私に忍び寄るのをやめてください、ありがとう

#include <bits/stdc++.h>
using namespace std;
#pragma comment(linker, "/stack:200000000")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")
#define ll long long
#define int long long
#define mod 1000000007 // 998244353
#define fast                                                                   \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(0);                                                                  \\
  cout.tie(0);
#define f(i, n) for (ll i = 0; i < n; i++)
#define fore(i, a, b) for (ll i = (ll)(a); i <= (ll)(b); ++i)
#define nl "\\n"
#define trace(x) cerr << #x << ": " << x << " " << endl
#define trace2(x, y) cerr << #x << ": " << x << " | " << #y << ": " << y << endl
#define trace3(x, y, z)                                                        \\
  cerr << #x << ":" << x << " | " << #y << ": " << y << " | " << #z << ": "    \\
       << z << endl
#define trace4(a, b, c, d)                                                     \\
  cerr << #a << ": " << a << " | " << #b << ": " << b << " | " << #c << ": "   \\
       << c << " | " << #d << ": " << d << endl
#define x first
#define y second
#define vc vector
#define pb push_back
#define ar array
#define all(a) (a).begin(), (a).end()
#define ms(v, n, x) fill(v, v + n, x);
#define init(c, a) memset(c, a, sizeof(c))
#define pll pair<ll, ll>
#define mll map<ll, ll>
#define sll set<ll>
#define vll vector<ll>
#define vpll vector<pll>
#define inf 9e18
#define cases                                                                  \\
  ll T;                                                                        \\
  cin >> T;                                                                    \\
  while (T--)
#define BLOCK 500
// const double PI = 3.141592653589793238460;
template <typename T> bool mmax(T &m, const T q) {
  if (m < q) {
    m = q;
    return true;
  } else
    return false;
}
template <typename T> bool mmin(T &m, const T q) {
  if (m > q) {
    m = q;
    return true;
  } else
    return false;
}
typedef std::complex<double> Complex;
typedef std::valarray<Complex> CArray;
void __print(int x) { cerr << x; }
void __print(long x) { cerr << x; }
// void __print(long long x) {cerr << x;}
void __print(unsigned x) { cerr << x; }
void __print(unsigned long x) { cerr << x; }
void __print(unsigned long long x) { cerr << x; }
void __print(float x) { cerr << x; }
void __print(double x) { cerr << x; }
void __print(long double x) { cerr << x; }
void __print(char x) { cerr << '\\'' << x << '\\''; }
void __print(const char *x) { cerr << '\\"' << x << '\\"'; }
void __print(const string &x) { cerr << '\\"' << x << '\\"'; }
void __print(bool x) { cerr << (x ? "true" : "false"); }

template <typename T, typename V> void __print(const pair<T, V> &x) {
  cerr << '{';
  __print(x.first);
  cerr << ',';
  __print(x.second);
  cerr << '}';
}
template <typename T> void __print(const T &x) {
  int f = 0;
  cerr << '{';
  for (auto &i : x)
    cerr << (f++ ? "," : ""), __print(i);
  cerr << "}";
}
void _print() { cerr << "]\\n"; }
template <typename T, typename... V> void _print(T t, V... v) {
  __print(t);
  if (sizeof...(v))
    cerr << ", ";
  _print(v...);
}
#ifndef ONLINE_JUDGE
#define debug(x...)                                                            \\
  cerr << "[" << #x << "] = [";                                                \\
  _print(x)
#else
#define debug(x...)
#endif

mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());

const int N = 1 << 20;

void add(ll &a, ll b) {
  a += b;
  if (a >= mod)
    a -= mod;
}

ll mul(ll a, ll b) { return (a * b) % mod; }

int power(int a, int b) {
  int res = 1;
  while (b) {
    if (b & 1) {
      res = res * a;
      res %= mod;
    }
    a = a * a;
    a %= mod;
    b = b / 2;
  }
  return res;
}

int inverse(int x) { return power(x, mod - 2); }

int32_t main() {

  fast cout << fixed << setprecision(12);

  ll n;
  cin >> n;

  vll a(n);
  ll sum = 0, si = 0;
  f(i, n) {
    cin >> a[i];
    add(sum, a[i]);
    add(si, mul(a[i], a[i]));
  }

  sum = mul(sum, sum);

  // debug(sum,si);

  ll ans = sum - si;
  if (ans < 0)
    ans += mod;

  ll t = ans * inverse(2) % mod;

  cout << t;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using ld = long double;
using cd = complex<ld>;

const ll MOD = 2e5 + 3;
const ll P_ROOT = 2; // primitive root
const ld PI = acos(-1);

ll pw[MOD + 5], ipw[MOD + 5];

void fft(vector<cd> &a, bool invert) {
  int n = (int)a.size();
  for (int i = 1, j = 0; i < n; i++) {
    int bit = n >> 1;
    for (; j & bit; bit >>= 1) {
      j ^= bit;
    }
    j ^= bit;
    if (i < j) {
      swap(a[i], a[j]);
    }
  }
  for (int len = 2; len <= n; len <<= 1) {
    ld 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 = a[i + j], v = a[i + j + len / 2] * w;
        a[i + j] = u + v;
        a[i + j + len / 2] = u - v;
        w *= wlen;
      }
    }
  }
  if (invert) {
    for (auto &x : a) {
      x /= n;
    }
  }
}

vector<cd> conv(vector<cd> const &a, vector<cd> const &b) {
  vector<cd> fa(a.begin(), a.end()), fb(b.begin(), b.end());
  int n = 1;
  while (n < (int)a.size() + b.size()) {
    n <<= 1;
  }
  fa.resize(n);
  fb.resize(n);
  fft(fa, false);
  fft(fb, false);
  for (int i = 0; i < n; i++) {
    fa[i] *= fb[i];
  }
  fft(fa, true);
  return fa;
}

void init() {
  ll temp = 1;
  ipw[temp] = 0;
  pw[0] = temp;
  for (int i = 1; i < MOD - 1; i++) {
    temp = (temp * P_ROOT) % MOD;
    ipw[temp] = i;
    pw[i] = temp;
  }
}

void solve() {
  init();
  int n;
  cin >> n;
  vector<int> A(MOD, 0);
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    ll x;
    cin >> x;
    if (x == 0)
      continue;
    ans = (ans - (x * x) % MOD);
    A[ipw[x]]++;
  }
  vector<cd> Ac(MOD);
  for (int i = 0; i < MOD; i++) {
    Ac[i] = cd(A[i]);
  }
  vector<cd> Sq = conv(Ac, Ac);
  int sz = (int)Sq.size();
  for (int i = 0; i < MOD - 1; i++) {
    ll temp = 0;
    if (i < sz)
      temp += (ll)roundl(Sq[i].real());
    if (i + MOD - 1 < sz)
      temp += (ll)roundl(Sq[i + MOD - 1].real());
    ans = (ans + temp * (pw[i]));
  }
  ans /= 2;
  cout << ans << '\\n';
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  int TC = 1;
  for (int i = 1; i <= TC; i++) {
    solve();
  }

  return 0;
}