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

#include <bits/stdc++.h>
using namespace std;
#define modulo 200003
#define mod(mod_x) ((((long long)mod_x + modulo)) % modulo)
#define Inf 1000000001

int beki(long long a, long long b, int M = modulo) {
  int x = 1;
  while (b != 0) {
    if (b & 1) {
      x = ((long long)x * a) % M;
    }
    a = ((long long)a * a) % M;
    b >>= 1;
  }
  return x;
}

long long findPrimitiveRoot(long long n) {
  long long P = n;
  n--;
  vector<long long> p;
  for (long long i = 2; i * i <= n; i++) {
    if (n % i == 0) {
      p.push_back(i);

      while (n % i == 0)
        n /= i;
    }
  }

  if (n != 1)
    p.push_back(n);

  for (long long i = 2; true; i++) {
    bool f = true;
    for (int j = 0; j < p.size(); j++) {
      if (beki(i, (P - 1) / p[j], P) == 1) {
        f = false;
        break;
      }
    }
    if (f)
      return i;
  }
}

void FFT(vector<complex<double>> &A, int sign) {
  int n = A.size();

  double PI = acos(-1.0);

  vector<complex<double>> tmp(n);
  int mask = n - 1;

  for (int i = n >> 1; i >= 1; i >>= 1) {
    double dir = 2.0 * PI * i * sign / n;

    complex<double> w(cos(dir), sin(dir));

    complex<double> now = 1.0;

    for (int j = 0; j < n; j += i) {
      for (int k = 0; k < i; k++) {
        tmp[j + k] =
            A[((j << 1) & mask) | k] + now * A[(((j << 1) | i) & mask) | k];
      }
      now *= w;
    }
    swap(A, tmp);
  }

  int num;
  for (int i = 0; true; i++) {
    if ((n >> i) & 1) {
      num = i;
      break;
    }
  }

  if (num & 1) {
    swap(A, tmp);
    for (int i = 0; i < n; i++)
      A[i] = tmp[i];
  }

  return;
}

vector<long long> do_FFT(vector<long long> &A, vector<long long> &B) {

  int n = A.size() + B.size() - 1;

  if (min(A.size(), B.size()) <= 150) {
    vector<long long> ret(n, 0);
    for (int i = 0; i < A.size(); i++) {
      for (int j = 0; j < B.size(); j++) {
        ret[i + j] = mod(ret[i + j] + mod(A[i] * B[j]));
      }
    }
    return ret;
  }

  for (int i = 0; true; i++) {
    if ((1 << i) >= n) {
      n = (1 << i);
      break;
    }
  }

  vector<complex<double>> a(n, complex<double>(0.0, 0.0)),
      b(n, complex<double>(0.0, 0.0));

  for (int i = 0; i < A.size(); i++)
    a[i].real(A[i]);
  for (int i = 0; i < B.size(); i++)
    b[i].real(B[i]);

  FFT(a, 1);
  if (A == B)
    b = a;
  else
    FFT(b, 1);

  vector<complex<double>> c(n);
  for (int i = 0; i < n; i++)
    c[i] = a[i] * b[i];

  FFT(c, -1);

  vector<long long> C(n);
  for (int i = 0; i < n; i++) {
    C[i] = round(c[i].real() / n);
  }

  return C;
}

int main() {

  int g = findPrimitiveRoot(modulo);

  vector<long long> arr;
  vector<int> gyaku(modulo, 0);

  for (int i = 0; i < modulo - 1; i++) {
    arr.push_back(beki(g, i));
    gyaku[arr.back()] = i;
  }

  int N;
  cin >> N;
  vector<int> A;
  vector<long long> X(modulo, 0LL);

  for (int i = 0; i < N; i++) {
    int a;
    cin >> a;

    if (a != 0) {
      A.push_back(a);
      X[gyaku[a]]++;
    }
  }

  vector<long long> Y = X;
  X = do_FFT(X, Y);

  for (int i = 0; i < A.size(); i++) {
    X[(gyaku[A[i]] + gyaku[A[i]]) % (modulo - 1)]--;
  }
  vector<long long> t(modulo - 1);
  long long ans = 0LL;
  for (int i = 0; i < X.size(); i++) {
    t[i % (modulo - 1)] += X[i];
  }

  for (int i = 0; i < t.size(); i++) {
    t[i] /= 2;
    // if(t[i]!=0)cout<<i<<','<<t[i]<<endl;
    ans += t[i] * arr[i];
  }

  cout << ans << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define modulo 200003
#define mod(mod_x) ((((long long)mod_x + modulo)) % modulo)
#define Inf 1000000001

int main() {

  int N;
  cin >> N;
  vector<int> A;
  vector<long long> S(modulo, 0);
  vector<long long> C(modulo, 0);

  for (int i = 0; i < N; i++) {
    int a;
    cin >> a;

    if (a != 0) {
      A.push_back(a);
      C[a]++;
    }
  }

  for (int i = 1; i < modulo; i++) {
    S[i] = C[i];
    S[i] += S[i - 1];
  }

  vector<long long> sum(modulo, 0);

  for (int i = 1; i < modulo; i++) {
    sum[i] = C[i] * i;
    sum[i] += sum[i - 1];
  }

  long long ans = 0LL;

  for (long long i = 1; i < modulo; i++) {
    if (C[i] == 0)
      continue;

    for (long long j = 0; j < i; j++) {
      int l, r;
      l = (j * modulo + i - 1) / i;
      r = ((1 + j) * modulo + i - 1) / i;
      r--;
      r = min(r, modulo - 1);
      l = max(l, 1);
      // cout<<i<<','<<j<<','<<l<<','<<r<<endl;
      ans +=
          (i * (sum[r] - sum[l - 1]) - (j * modulo * (S[r] - S[l - 1]))) * C[i];
      // if((i * (sum[r]-sum[l-1]) - (j*modulo*(S[r]-S[l-1]))) *
      // C[i]!=0)cout<<(i * (sum[r]-sum[l-1]) - (j*modulo*(S[r]-S[l-1]))) *
      // C[i]<<endl;
    }
  }

  ans /= 2;
  long long t = 0LL;
  for (int i = 0; i < A.size(); i++) {
    // cout<<mod(A[i]*A[i])<<endl;
    t += mod(A[i] * A[i]);
  }

  cout << ans - t / 2 << endl;

  return 0;
}