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

#include <bits/stdc++.h>
using namespace std;

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()
#define pb push_back
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

/**
 * Author: Ludo Pulles, chilli, Simon Lindholm
 * Date: 2019-01-09
 * License: CC0
 * Source: http://neerc.ifmo.ru/trains/toulouse/2017/fft2.pdf (do read, it's
 excellent) Accuracy bound from http://www.daemonology.net/papers/fft.pdf
 * Description: fft(a) computes $\\hat f(k) = \\sum_x a[x] \\exp(2\\pi i \\cdot k x /
 N)$ for all $k$. N must be a power of 2. Useful for convolution:
   \\texttt{conv(a, b) = c}, where $c[x] = \\sum a[i]b[x-i]$.
   For convolution of complex numbers or more than two vectors: FFT, multiply
   pointwise, divide by n, reverse(start+1, end), FFT back.
   Rounding is safe if $(\\sum a_i^2 + \\sum b_i^2)\\log_2{N} < 9\\cdot10^{14}$
   (in practice $10^{16}$; higher for random inputs).
   Otherwise, use NTT/FFTMod.
 * Time: O(N \\log N) with $N = |A|+|B|$ ($\\tilde 1s$ for $N=2^{22}$)
 * Status: somewhat tested
 * Details: An in-depth examination of precision for both FFT and FFTMod can be
 found
 * here
 (https://github.com/simonlindholm/fft-precision/blob/master/fft-precision.md)
 */
#pragma once

typedef complex<double> C;
typedef vector<double> vd;
void fft(vector<C> &a) {
  int n = sz(a), L = 31 - __builtin_clz(n);
  static vector<complex<long double>> R(2, 1);
  static vector<C> rt(2, 1); // (^ 10% faster if double)
  for (static int k = 2; k < n; k *= 2) {
    R.resize(n);
    rt.resize(n);
    auto x = polar(1.0L, acos(-1.0L) / k);
    rep(i, k, 2 * k) rt[i] = R[i] = i & 1 ? R[i / 2] * x : R[i / 2];
  }
  vi rev(n);
  rep(i, 0, n) rev[i] = (rev[i / 2] | (i & 1) << L) / 2;
  rep(i, 0, n) if (i < rev[i]) swap(a[i], a[rev[i]]);
  for (int k = 1; k < n; k *= 2)
    for (int i = 0; i < n; i += 2 * k)
      rep(j, 0, k) {
        // C z = rt[j+k] * a[i+j+k]; // (25% faster if hand-rolled)  ///
        // include-line
        auto x = (double *)&rt[j + k],
             y = (double *)&a[i + j + k]; /// exclude-line
        C z(x[0] * y[0] - x[1] * y[1],
            x[0] * y[1] + x[1] * y[0]); /// exclude-line
        a[i + j + k] = a[i + j] - z;
        a[i + j] += z;
      }
}
vd conv(const vd &a, const vd &b) {
  if (a.empty() || b.empty())
    return {};
  vd res(sz(a) + sz(b) - 1);
  int L = 32 - __builtin_clz(sz(res)), n = 1 << L;
  vector<C> in(n), out(n);
  copy(all(a), begin(in));
  rep(i, 0, sz(b)) in[i].imag(b[i]);
  fft(in);
  for (C &x : in)
    x *= x;
  rep(i, 0, n) out[i] = in[-i & (n - 1)] - conj(in[i]);
  fft(out);
  rep(i, 0, sz(res)) res[i] = imag(out[i]) / (4 * n);
  return res;
}

int n;
ll p = 200003, cnt[200005], ans;
vd a, b;
int main() {
  cin.tie(0)->sync_with_stdio(0);
  cin.exceptions(cin.failbit);
  cin >> n;
  rep(i, 0, n) {
    int a;
    cin >> a;
    cnt[a]++;
  }
  ll x = 1;
  for (int i = 0; i < p - 1; i++) {
    a.pb(cnt[x]);
    b.pb(cnt[x]);
    x = x * 2 % p;
  }
  vd c = conv(a, b);
  x = 1;
  for (int i = 0; i < sz(c); i++) {
    ans += x * ll(round(c[i]));
    x = x * 2 % p;
  }
  for (ll i = 1; i < p; i++) {
    ans -= ((i * i) % p) * cnt[i];
  }
  cout << ans / 2LL << endl;
}

Error: TLE

Fixed submission:

#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
#define rng(i, l, r) for (int i = int(l); i < int(r); i++)
#define rep(i, r) rng(i, 0, r)
#define rrng(i, l, r) for (int i = int(r) - 1; i >= int(l); i--)
#define rrep(i, r) rrng(i, 0, r)
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define mt make_tuple
#define F first
#define S second
#define bg begin()
#define ed end()
#define all(x) x.bg, x.ed
#define si(x) int(x.size())
#define inf INT_MAX / 2 - 100
#define infl LLONG_MAX / 3
#ifdef LOCAL
#define dmp(x) cerr << __LINE__ << ' ' << #x << ' ' << x << endl
#else
#define dmp(x) void(0)
#endif

template <class t, class u> bool chmax(t &a, u b) {
  if (a < b)
    a = b;
  return a < b;
}
template <class t, class u> bool chmin(t &a, u b) {
  if (b < a)
    a = b;
  return b < a;
}

template <class t> using vc = vector<t>;
template <class t> using vvc = vector<vector<t>>;

using pi = pair<int, int>;
using pl = pair<ll, ll>;
using vi = vc<int>;
using vl = vc<ll>;

ll readl(void) {
  ll x;
  cin >> x;
  return x;
}
int readi(void) {
  int x;
  cin >> x;
  return x;
}

string readstr() {
  string s;
  cin >> s;
  return s;
}
vi readvi(int n, int off = 0) {
  vi v(n);
  rep(i, n) v[i] = readi(), v[i] += off;
  return v;
}

vl readvl(int n, int off = 0) {
  vl v(n);
  rep(i, n) v[i] = readl(), v[i] += off;
  return v;
}

template <class t> void print(t x, int suc = 1) {
  cout << x;
  if (suc == 1)
    cout << "\\n";
  if (suc == 2)
    cout << " ";
}

template <class t> void print(const vc<t> &v, int suc = 1) {
  rep(i, si(v)) print(v[i], i == int(si(v)) - 1 ? 1 : suc);
}

template <class t> bool inc(t a, t b, t c) { return !(c < b || b < a); }

template <class t> void compress(vc<t> &v) {
  sort(all(v));
  v.erase(unique(all(v)), v.ed);
}

template <class t> int lwb(const vc<t> &v, const t &a) {
  return lower_bound(all(v), a) - v.bg;
}

template <class t> struct Compress {
  vc<t> v;
  Compress() = default;
  Compress(const vc<t> &x) { add(x); }
  Compress(const initializer_list<vc<t>> &x) {
    for (auto &p : x)
      add(p);
  }
  void add(const t &x) { v.eb(x); }
  void add(const vc<t> &x) { copy(all(x), back_inserter(v)); }
  void build() { compress(v); }
  int get(const t &x) const { return lwb(v, x); }
  vc<t> get(const vc<t> &x) const {
    vc<t> res(x);
    for (auto &p : res)
      p = get(p);
    return res;
  }
  const t &operator[](int x) const { return v[x]; }
  int size() { return v.size(); }
};
void Yes(bool ex = true) {
  cout << "Yes\\n";
  if (ex)
    exit(0);
}
void YES(bool ex = true) {
  cout << "YES\\n";
  if (ex)
    exit(0);
}
void No(bool ex = true) {
  cout << "No\\n";
  if (ex)
    exit(0);
}
void NO(bool ex = true) {
  cout << "NO\\n";
  if (ex)
    exit(0);
}
void orYes(bool x, bool ex = true) {
  if (x)
    Yes(ex);
  else
    No(ex);
}
void orYES(bool x, bool ex = true) {
  if (x)
    YES(ex);
  else
    NO(ex);
}
void Possible(bool ex = true) {
  cout << "Possible\\n";
  if (ex)
    exit(0);
}
void POSSIBLE(bool ex = true) {
  cout << "POSSIBLE\\n";
  if (ex)
    exit(0);
}
void Impossible(bool ex = true) {
  cout << "Impossible\\n";
  if (ex)
    exit(0);
}
void IMPOSSIBLE(bool ex = true) {
  cout << "IMPOSSIBLE\\n";
  if (ex)
    exit(0);
}
void orPossible(bool x, bool ex = true) {
  if (x)
    Possible(ex);
  else
    Impossible(ex);
}
void orPOSSIBLE(bool x, bool ex = true) {
  if (x)
    POSSIBLE(ex);
  else
    IMPOSSIBLE(ex);
}

using uint = unsigned;
using ull = unsigned long long;

template <uint const &MOD> struct Modular {
  static constexpr uint const &mod = MOD;
  uint v;
  Modular(long long x = 0) { c(x % mod + mod); }
  Modular &c(uint x) {
    v = x < mod ? x : x - mod;
    return *this;
  }
  Modular pow(int k) const {
    Modular res(1), tmp(v);
    while (k) {
      if (k & 1)
        res *= tmp;
      tmp *= tmp;
      k >>= 1;
    }
    return res;
  }

  Modular inv() const { return pow(mod - 2); }
  Modular operator-() const { return Modular(mod - v); }

  Modular &operator+=(const Modular &x) { return c(v + x.v); }
  Modular &operator-=(const Modular &x) { return c(v + mod - x.v); }
  Modular &operator*=(const Modular &x) {
    v = ull(v) * x.v % mod;
    return *this;
  }
  Modular &operator/=(const Modular &x) { return *this *= x.inv(); }

  Modular operator+(const Modular &x) const { return Modular(*this) += x; }
  Modular operator-(const Modular &x) const { return Modular(*this) -= x; }
  Modular operator*(const Modular &x) const { return Modular(*this) *= x; }
  Modular operator/(const Modular &x) const { return Modular(*this) /= x; }

  friend Modular operator+(long long x, const Modular &y) {
    return Modular(x) + y;
  }
  friend Modular operator-(long long x, const Modular &y) {
    return Modular(x) - y;
  }
  friend Modular operator*(long long x, const Modular &y) {
    return Modular(x) * y;
  }
  friend Modular operator/(long long x, const Modular &y) {
    return Modular(x) / y;
  }

  friend ostream &operator<<(ostream &os, const Modular &x) {
    return os << x.v;
  }
  friend istream &operator>>(istream &is, Modular &x) {
    long long p;
    is >> p;
    x = Modular(p);
    return is;
  }
  bool operator==(const Modular &x) const { return v == x.v; }
  bool operator!=(const Modular &x) const { return v != x.v; }
  bool operator<(const Modular &x) const { return v < x.v; }
  explicit operator bool() const { return v; }
};

// uint MODULAR=998244353;
uint MODULAR = 1000000007;
using Mint = Modular<MODULAR>;

vector<Mint> fact, finv, invs;
void Initfact(int n = (1 << 21) + 10) {
  fact.resize(n + 1), finv.resize(n + 1), invs.resize(n + 1);
  fact[0] = 1;
  for (int i = 1; i <= n; i++) {
    fact[i] = fact[i - 1] * i;
  }
  finv[n] = fact[n].inv();
  for (int i = n - 1; i >= 0; i--) {
    finv[i] = finv[i + 1] * (i + 1);
  }
  invs[0] = 1;
  for (int i = 0; i <= n; i++) {
    invs[i] = finv[i] * fact[i - 1];
  }
}

Mint comb(int n, int k) { return fact[n] * finv[n - k] * finv[k]; }

int n;
ll ans, a[200005], b[200005];
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(0);
  cin >> n;
  for (int i = 0; i < n; i++)
    cin >> a[i];
  for (int i = 0; i < n; i++)
    b[i] = a[i];
  for (int i = 0; i < n; i++) {
    for (int j = i + 1; j < n; j++) {
      ans += a[i] * b[j] % 200003;
    }
  }
  print(ans);
}