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

// #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")
// // codeforces
// #pragma GCC target("avx,avx2,fma")
// #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,tune=native") //
// yandex

#include "bits/stdc++.h"
using namespace std;
#ifndef LOCAL
#define endl '\\n'
#endif

#define fr(i, a, b) for (int i = a; i <= b; i++)
#define rf(i, a, b) for (int i = a; i >= b; i--)
#define pf push_front
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define sz(x) (int)x.size()
#define lb lower_bound
#define ub upper_bound

typedef long long ll;
typedef double f80;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

int pct(int x) { return __builtin_popcount(x); }
int pct(ll x) { return __builtin_popcountll(x); }
int bt(int x) { return 31 - __builtin_clz(x); }  // floor(log2(x))
int bt(ll x) { return 63 - __builtin_clzll(x); } // floor(log2(x))
int cdiv(int a, int b) { return a / b + !(a < 0 || a % b == 0); }
ll cdiv(ll a, ll b) { return a / b + !(a < 0 || a % b == 0); }
int nxt_C(int x) {
  int c = x & -x, r = x + c;
  return (((r ^ x) >> 2) / c) | r;
}
ll nxt_C(ll x) {
  ll c = x & -x, r = x + c;
  return (((r ^ x) >> 2) / c) | r;
}

vector<int> get_bits(int mask) {
  vector<int> bb;
  while (mask) {
    int b = bt(mask);
    bb.pb(b);
    mask ^= (1 << b);
  }
  reverse(all(bb));
  return bb;
}

int get_mask(vector<int> v) {
  int mask = 0;
  for (int x : v) {
    mask ^= (1 << x);
  }
  return mask;
}

template <typename T> void uniq(vector<T> &v) {
  sort(all(v));
  v.resize(unique(all(v)) - v.begin());
}

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

ll rand(ll l, ll r) {
  uniform_int_distribution<ll> uid(l, r);
  return uid(rng);
}

void sc() {}

template <typename Head, typename... Tail> void sc(Head &H, Tail &...T) {
  cin >> H;
  sc(T...);
}

#ifdef LOCAL
#define debug(...)                                                             \\
  cerr << "[L:" << __LINE__ << "][" << #__VA_ARGS__ << "]:",                   \\
      debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

// #ifndef LOCAL
// string to_string(__int128 x) {
// 	string s = "";
// 	bool neg = 0;
// 	if(x < 0) { s += "-"; neg = 1; x = -x; }
// 	if(!x) s += '0';
// 	while(x) {
// 		int rem = x % 10;
// 		s += to_string(rem);
// 		x /= 10;
// 	}
// 	reverse(s.begin() + neg, s.end());
// 	return s;
// }
// #endif

const int mod = 1e9 + 7; // 998244353;

int pwr(int a, ll b, int mod) {
  int ans = 1;
  while (b) {
    if (b & 1)
      ans = (ans * 1LL * a) % mod;
    a = (a * 1LL * a) % mod;
    b >>= 1;
  }
  return ans;
}

/*
        Lookout for overflows!!
        Check array sizes!!
        Clear before test cases!!
        Use the correct modulo!!
        Check for corner cases!!
        Are you forgetting something?!
        Read problem statement carefully!!!
*/

struct base {
  f80 x, y;
  base() { x = y = 0; }
  base(f80 _x, f80 _y) { x = _x, y = _y; }
  base(f80 _x) { x = _x, y = 0; }
  void operator=(f80 _x) { x = _x, y = 0; }
  f80 real() { return x; }
  f80 imag() { return y; }
  base operator+(const base &b) { return base(x + b.x, y + b.y); }
  void operator+=(const base &b) { x += b.x, y += b.y; }
  base operator*(const base &b) {
    return base(x * b.x - y * b.y, x * b.y + y * b.x);
  }
  void operator*=(const base &b) {
    f80 p = x * b.x - y * b.y, q = x * b.y + y * b.x;
    x = p, y = q;
  }
  void operator/=(f80 k) { x /= k, y /= k; }
  base operator-(const base &b) { return base(x - b.x, y - b.y); }
  void operator-=(const base &b) { x -= b.x, y -= b.y; }
  base conj() { return base(x, -y); }
  base operator/(f80 k) { return base(x / k, y / k); }
  void Print() { cerr << x << " + " << y << "i\\n"; }
};
double PI = 2.0 * acos(0.0);
const int MAXN = 19;
const int maxn = (1 << MAXN);
base W[maxn], invW[maxn], P1[maxn], Q1[maxn];
void precompute_powers() {
  for (int i = 0; i < maxn / 2; i++) {
    double ang = (2 * PI * i) / maxn;
    f80 _cos = cos(ang), _sin = sin(ang);
    W[i] = base(_cos, _sin);
    invW[i] = base(_cos, -_sin);
  }
}
void fft(vector<base> &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) {
    for (int i = 0; i < n; i += len) {
      int ind = 0, add = maxn / len;
      for (int j = 0; j < len / 2; ++j) {
        base u = a[i + j],
             v = (a[i + j + len / 2] * (invert ? invW[ind] : W[ind]));
        a[i + j] = (u + v);
        a[i + j + len / 2] = (u - v);
        ind += add;
      }
    }
  }
  if (invert)
    for (int i = 0; i < n; ++i)
      a[i] /= n;
}

// 4 FFTs in total for a precise convolution
void mul_big_mod(vector<ll> &a, vector<ll> &b, ll mod) {
  int n1 = a.size(), n2 = b.size();
  int final_size = a.size() + b.size() - 1;
  int n = 1;
  while (n < final_size)
    n <<= 1;
  vector<base> P(n), Q(n);
  int SQRTMOD = (int)sqrt(mod) + 10;
  for (int i = 0; i < n1; i++)
    P[i] = base(a[i] % SQRTMOD, a[i] / SQRTMOD);
  for (int i = 0; i < n2; i++)
    Q[i] = base(b[i] % SQRTMOD, b[i] / SQRTMOD);
  fft(P, 0);
  fft(Q, 0);
  base A1, A2, B1, B2, X, Y;
  for (int i = 0; i < n; i++) {
    X = P[i];
    Y = P[(n - i) & (n - 1)].conj();
    A1 = (X + Y) * base(0.5, 0);
    A2 = (X - Y) * base(0, -0.5);
    X = Q[i];
    Y = Q[(n - i) & (n - 1)].conj();
    B1 = (X + Y) * base(0.5, 0);
    B2 = (X - Y) * base(0, -0.5);
    P1[i] = A1 * B1 + A2 * B2 * base(0, 1);
    Q1[i] = A1 * B2 + A2 * B1;
  }
  for (int i = 0; i < n; i++)
    P[i] = P1[i], Q[i] = Q1[i];
  fft(P, 1);
  fft(Q, 1);
  a.resize(final_size);
  for (int i = 0; i < final_size; i++) {
    ll x = (ll)(P[i].real() + 0.5);
    ll y = (ll)(P[i].imag() + 0.5) % mod;
    ll z = (ll)(Q[i].real() + 0.5);
    a[i] = (x + ((y * SQRTMOD + z) % mod) * SQRTMOD) % mod;
  }
}
vector<ll> sq(const vector<ll> &a) {
  int tot = a.size() + a.size() - 1;
  int L = 32 - __builtin_clz(tot), n = 1 << L;
  vector<base> c(n, 0);
  for (int i = 0; i < a.size(); i++) {
    c[i].x = a[i];
    c[i].y = a[i];
  }
  fft(c, 0);
  vector<base> d(n, 0);
  for (int i = 0; i < n; i++) {
    base A = c[i], B = c[(n - i) & (n - 1)].conj();
    base p = (A + B) * base(0.5, 0);
    base q = (A - B) * base(0, -0.5);
    d[i] = p * q;
  }
  fft(d, 1);
  vector<ll> ans(tot);
  for (int i = 0; i < tot; i++) {
    ans[i] = round(d[i].x);
  }
  return ans;
}
// vector<int> v1, v2;
// void solve(){
//     int n, m;
//     cin >> n >> m;
//     v1.resize(n + 1), v2.resize(m + 1);
//     fr(i, 0, n)
//         cin >> v1[i];
//     fr(i, 0, m)
//         cin >> v2[i];
//     mul_big_mod(v1, v2, mod);
//     fr(i, 0, n + m)
//         cout << v1[i] << " ";
//     cout << endl;
// }

// Returns minimum x for which a ^ x % m = b % m.
const int N = 2e5 + 5;
const int P = 2e5 + 3;

int a[N];
int g[P];

vector<ll> v(P, 0);

void solve() {
  precompute_powers();
  g[1] = 0;
  int xx = 1;
  for (int i = 1; i < P; i++) {
    xx = (xx * 2) % P;
    g[xx] = i;
  }
  int n;
  sc(n);
  fr(i, 1, n) {
    sc(a[i]);
    if (!a[i])
      continue;
    int k = g[a[i]];
    v[k]++;
  }
  v = sq(v);
  ll x = 1;
  ll ans = 0;
  for (int i = 0; i < sz(v); i++) {
    ans += v[i] * x;
    x *= 2;
    x %= P;
  }
  fr(i, 1, n) { ans -= ((a[i] * 1LL * a[i]) % P); }
  ans /= 2;
  cout << ans;
}

int main() {
  ios ::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  int t = 1;
  // cin >> t;
  for (int tt = 1; tt <= t; tt++) {
    solve();
  }
  return 0;
}

Error: TLE

Fixed submission:

// #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")
// // codeforces
// #pragma GCC target("avx,avx2,fma")
// #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,tune=native") //
// yandex

#include "bits/stdc++.h"
using namespace std;
#ifndef LOCAL
#define endl '\\n'
#endif

#define fr(i, a, b) for (int i = a; i <= b; i++)
#define rf(i, a, b) for (int i = a; i >= b; i--)
#define pf push_front
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define sz(x) (int)x.size()
#define lb lower_bound
#define ub upper_bound

typedef long long ll;
typedef long double f80;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

int pct(int x) { return __builtin_popcount(x); }
int pct(ll x) { return __builtin_popcountll(x); }
int bt(int x) { return 31 - __builtin_clz(x); }  // floor(log2(x))
int bt(ll x) { return 63 - __builtin_clzll(x); } // floor(log2(x))
int cdiv(int a, int b) { return a / b + !(a < 0 || a % b == 0); }
ll cdiv(ll a, ll b) { return a / b + !(a < 0 || a % b == 0); }
int nxt_C(int x) {
  int c = x & -x, r = x + c;
  return (((r ^ x) >> 2) / c) | r;
}
ll nxt_C(ll x) {
  ll c = x & -x, r = x + c;
  return (((r ^ x) >> 2) / c) | r;
}

vector<int> get_bits(int mask) {
  vector<int> bb;
  while (mask) {
    int b = bt(mask);
    bb.pb(b);
    mask ^= (1 << b);
  }
  reverse(all(bb));
  return bb;
}

int get_mask(vector<int> v) {
  int mask = 0;
  for (int x : v) {
    mask ^= (1 << x);
  }
  return mask;
}

template <typename T> void uniq(vector<T> &v) {
  sort(all(v));
  v.resize(unique(all(v)) - v.begin());
}

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

ll rand(ll l, ll r) {
  uniform_int_distribution<ll> uid(l, r);
  return uid(rng);
}

void sc() {}

template <typename Head, typename... Tail> void sc(Head &H, Tail &...T) {
  cin >> H;
  sc(T...);
}

#ifdef LOCAL
#define debug(...)                                                             \\
  cerr << "[L:" << __LINE__ << "][" << #__VA_ARGS__ << "]:",                   \\
      debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

// #ifndef LOCAL
// string to_string(__int128 x) {
// 	string s = "";
// 	bool neg = 0;
// 	if(x < 0) { s += "-"; neg = 1; x = -x; }
// 	if(!x) s += '0';
// 	while(x) {
// 		int rem = x % 10;
// 		s += to_string(rem);
// 		x /= 10;
// 	}
// 	reverse(s.begin() + neg, s.end());
// 	return s;
// }
// #endif

const int mod = 1e9 + 7; // 998244353;

int pwr(int a, ll b, int mod) {
  int ans = 1;
  while (b) {
    if (b & 1)
      ans = (ans * 1LL * a) % mod;
    a = (a * 1LL * a) % mod;
    b >>= 1;
  }
  return ans;
}

/*
        Lookout for overflows!!
        Check array sizes!!
        Clear before test cases!!
        Use the correct modulo!!
        Check for corner cases!!
        Are you forgetting something?!
        Read problem statement carefully!!!
*/

struct base {
  f80 x, y;
  base() { x = y = 0; }
  base(f80 _x, f80 _y) { x = _x, y = _y; }
  base(f80 _x) { x = _x, y = 0; }
  void operator=(f80 _x) { x = _x, y = 0; }
  f80 real() { return x; }
  f80 imag() { return y; }
  base operator+(const base &b) { return base(x + b.x, y + b.y); }
  void operator+=(const base &b) { x += b.x, y += b.y; }
  base operator*(const base &b) {
    return base(x * b.x - y * b.y, x * b.y + y * b.x);
  }
  void operator*=(const base &b) {
    f80 p = x * b.x - y * b.y, q = x * b.y + y * b.x;
    x = p, y = q;
  }
  void operator/=(f80 k) { x /= k, y /= k; }
  base operator-(const base &b) { return base(x - b.x, y - b.y); }
  void operator-=(const base &b) { x -= b.x, y -= b.y; }
  base conj() { return base(x, -y); }
  base operator/(f80 k) { return base(x / k, y / k); }
  void Print() { cerr << x << " + " << y << "i\\n"; }
};
double PI = 2.0 * acos(0.0);
const int MAXN = 19;
const int maxn = (1 << MAXN);
base W[maxn], invW[maxn], P1[maxn], Q1[maxn];
void precompute_powers() {
  for (int i = 0; i < maxn / 2; i++) {
    double ang = (2 * PI * i) / maxn;
    f80 _cos = cos(ang), _sin = sin(ang);
    W[i] = base(_cos, _sin);
    invW[i] = base(_cos, -_sin);
  }
}
void fft(vector<base> &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) {
    for (int i = 0; i < n; i += len) {
      int ind = 0, add = maxn / len;
      for (int j = 0; j < len / 2; ++j) {
        base u = a[i + j],
             v = (a[i + j + len / 2] * (invert ? invW[ind] : W[ind]));
        a[i + j] = (u + v);
        a[i + j + len / 2] = (u - v);
        ind += add;
      }
    }
  }
  if (invert)
    for (int i = 0; i < n; ++i)
      a[i] /= n;
}

// 4 FFTs in total for a precise convolution
void mul_big_mod(vector<ll> &a, vector<ll> &b, ll mod) {
  int n1 = a.size(), n2 = b.size();
  int final_size = a.size() + b.size() - 1;
  int n = 1;
  while (n < final_size)
    n <<= 1;
  vector<base> P(n), Q(n);
  int SQRTMOD = (int)sqrt(mod) + 10;
  for (int i = 0; i < n1; i++)
    P[i] = base(a[i] % SQRTMOD, a[i] / SQRTMOD);
  for (int i = 0; i < n2; i++)
    Q[i] = base(b[i] % SQRTMOD, b[i] / SQRTMOD);
  fft(P, 0);
  fft(Q, 0);
  base A1, A2, B1, B2, X, Y;
  for (int i = 0; i < n; i++) {
    X = P[i];
    Y = P[(n - i) & (n - 1)].conj();
    A1 = (X + Y) * base(0.5, 0);
    A2 = (X - Y) * base(0, -0.5);
    X = Q[i];
    Y = Q[(n - i) & (n - 1)].conj();
    B1 = (X + Y) * base(0.5, 0);
    B2 = (X - Y) * base(0, -0.5);
    P1[i] = A1 * B1 + A2 * B2 * base(0, 1);
    Q1[i] = A1 * B2 + A2 * B1;
  }
  for (int i = 0; i < n; i++)
    P[i] = P1[i], Q[i] = Q1[i];
  fft(P, 1);
  fft(Q, 1);
  a.resize(final_size);
  for (int i = 0; i < final_size; i++) {
    ll x = (ll)(P[i].real() + 0.5);
    ll y = (ll)(P[i].imag() + 0.5) % mod;
    ll z = (ll)(Q[i].real() + 0.5);
    a[i] = (x + ((y * SQRTMOD + z) % mod) * SQRTMOD) % mod;
  }
}
vector<ll> mul(const vector<ll> &a, const vector<ll> &b) {
  int tot = a.size() + b.size() - 1;
  int L = 32 - __builtin_clz(tot), n = 1 << L;
  vector<base> c(n, 0);
  for (int i = 0; i < a.size(); i++) {
    c[i].x = a[i];
  }
  for (int i = 0; i < b.size(); i++) {
    c[i].y = b[i];
  }
  fft(c, 0);
  vector<base> d(n, 0);
  for (int i = 0; i < n; i++) {
    base A = c[i], B = c[(n - i) & (n - 1)].conj();
    base p = (A + B) * base(0.5, 0);
    base q = (A - B) * base(0, -0.5);
    d[i] = p * q;
  }
  fft(d, 1);
  vector<ll> ans(tot);
  for (int i = 0; i < tot; i++) {
    ans[i] = round(d[i].x);
  }
  return ans;
}
// vector<int> v1, v2;
// void solve(){
//     int n, m;
//     cin >> n >> m;
//     v1.resize(n + 1), v2.resize(m + 1);
//     fr(i, 0, n)
//         cin >> v1[i];
//     fr(i, 0, m)
//         cin >> v2[i];
//     mul_big_mod(v1, v2, mod);
//     fr(i, 0, n + m)
//         cout << v1[i] << " ";
//     cout << endl;
// }

// Returns minimum x for which a ^ x % m = b % m.
int solve(int a, int b, int m) {
  a %= m, b %= m;
  int k = 1, add = 0, g;
  while ((g = __gcd(a, m)) > 1) {
    if (b == k)
      return add;
    if (b % g)
      return -1;
    b /= g, m /= g, ++add;
    k = (k * 1ll * a / g) % m;
  }

  int n = sqrt(m) + 1;
  int an = 1;
  for (int i = 0; i < n; ++i)
    an = (an * 1ll * a) % m;

  unordered_map<int, int> vals;
  for (int q = 0, cur = b; q <= n; ++q) {
    vals[cur] = q;
    cur = (cur * 1ll * a) % m;
  }

  for (int p = 1, cur = k; p <= n; ++p) {
    cur = (cur * 1ll * an) % m;
    if (vals.count(cur)) {
      int ans = n * p - vals[cur] + add;
      return ans;
    }
  }
  return -1;
}

const int N = 2e5 + 5;
const int P = 2e5 + 3;
int a[N];

vector<ll> v(P, 0);

void solve() {
  precompute_powers();
  int n;
  sc(n);
  fr(i, 1, n) {
    sc(a[i]);
    if (!a[i])
      continue;
    int k = solve(2, a[i], P);
    debug(k);
    v[k]++;
  }
  v = mul(v, v);
  ll x = 1;
  ll ans = 0;
  for (int i = 0; i < sz(v); i++) {
    if (v[i])
      debug(i, v[i]);
    ans += v[i] * x;
    x *= 2;
    x %= P;
  }
  fr(i, 1, n) { ans -= ((a[i] * 1LL * a[i]) % P); }
  ans /= 2;
  cout << ans;
}

int main() {
  ios ::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  int t = 1;
  // cin >> t;
  for (int tt = 1; tt <= t; tt++) {
    solve();
  }
  return 0;
}