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

#include <bits/stdc++.h>

using namespace std;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < v.size(); ++i) {
    if (i) {
      os << " ";
    }
    os << v[i];
  }
  os << "]";
  return os;
}

const int max_lev = 19, mod_fft = 998244353;

template <int mod, int max_lev> struct NTT {
  static const int max_n = 1 << max_lev;

  int roots[2][max_lev];
  int rev[max_n], buf1[max_n], buf2[max_n], bufpw[max_n];

  NTT() {
    int root = find_root();
    int rroot = power(root, mod - 2);
    roots[0][0] = root;
    roots[1][0] = rroot;
    for (int i = 1; i < max_lev; ++i) {
      roots[0][i] = mul(roots[0][i - 1], roots[0][i - 1]);
      roots[1][i] = mul(roots[1][i - 1], roots[1][i - 1]);
    }
    for (int i = 1; i < max_n; ++i) {
      rev[i] = (rev[i / 2] / 2) + ((i & 1) << (max_lev - 1));
    }
  }

  static inline void inc(int &x, int y) {
    x += y;
    if (x >= mod) {
      x -= mod;
    }
  }

  static inline int mul(int x, int y) { return (1LL * x * y) % mod; }

  static int power(int x, int y) {
    if (y == 0) {
      return 1;
    }
    if (y % 2 == 0) {
      return power(mul(x, x), y / 2);
    }
    return mul(x, power(x, y - 1));
  }

  static int find_root() {
    for (int root = 2;; ++root) {
      if (power(root, max_n) == 1) {
        int cur = 1;
        bool ok = true;
        for (int i = 1; i < max_n; ++i) {
          cur = mul(cur, root);
          if (cur == 1) {
            ok = false;
            break;
          }
        }
        if (ok) {
          return root;
        }
      }
    }
  }

  void fft(int *a, const int *roots, int lg) {
    const int n = 1 << lg;
    for (int i = 0; i < n; ++i) {
      int r = (rev[i] >> (max_lev - lg)) & (n - 1);
      if (i < r) {
        swap(a[i], a[r]);
      }
    }
    int pos = max_lev - 1;
    for (int st = 1; st < n; st *= 2) {
      int cur = 1;
      for (int i = 0; i < st; ++i) {
        bufpw[i] = cur;
        cur = mul(cur, roots[pos]);
      }
      --pos;
      for (int i = 0; i < n; i += 2 * st) {
        int *curpw = bufpw;
        for (int j = i; j < i + st; ++j) {
          int y = mul(*curpw++, a[j + st]);
          a[j + st] = a[j];
          inc(a[j + st], mod - y);
          inc(a[j], y);
        }
      }
    }
  }

  vector<int> product(const vector<int> &a, const vector<int> &b) {
    int sz = a.size() + b.size() - 1;
    int n = 1, rn = 1, lg = 0;
    while (n < sz) {
      n *= 2;
      rn = mul(rn, (mod + 1) / 2);
      ++lg;
    }
    memcpy(buf1, a.data(), 4 * a.size());
    memset(buf1 + a.size(), 0, 4 * (n - a.size()));
    memcpy(buf2, b.data(), 4 * b.size());
    memset(buf2 + b.size(), 0, 4 * (n - b.size()));
    fft(buf1, roots[0], lg);
    fft(buf2, roots[0], lg);
    for (int i = 0; i < n; ++i) {
      buf1[i] = mul(buf1[i], buf2[i]);
    }
    fft(buf1, roots[1], lg);
    vector<int> ans(sz);
    for (int i = 0; i < sz; ++i) {
      ans[i] = mul(buf1[i], rn);
    }
    return ans;
  }

  vector<int> square(const vector<int> &a) {
    int sz = a.size() + a.size() - 1;
    int n = 1, rn = 1, lg = 0;
    while (n < sz) {
      n *= 2;
      rn = mul(rn, (mod + 1) / 2);
      ++lg;
    }
    memcpy(buf1, a.data(), 4 * a.size());
    memset(buf1 + a.size(), 0, 4 * (n - a.size()));
    fft(buf1, roots[0], lg);
    for (int i = 0; i < n; ++i) {
      buf1[i] = mul(buf1[i], buf1[i]);
    }
    fft(buf1, roots[1], lg);
    vector<int> ans(sz);
    for (int i = 0; i < sz; ++i) {
      ans[i] = mul(buf1[i], rn);
    }
    return ans;
  }
};

NTT<mod_fft, max_lev> ntt;

const int max_n = 200222, inf = 1000111222;
const int mod = 200003;
const int g = 2;

int mul(int x, int y) { return (1LL * x * y) % mod; }

int n, a[max_n];
int pw[max_n], lg[max_n];
long long ans;

const int sq = 555;

vector<int> A(mod), B(mod), C(mod);

int main() {
  // freopen("input.txt", "r", stdin);
  // freopen("output.txt", "w", stdout);
  pw[0] = 1;
  for (int i = 1; i + 1 < mod; ++i) {
    pw[i] = mul(pw[i - 1], g);
  }
  memset(lg, -1, sizeof(lg));
  for (int i = 0; i + 1 < mod; ++i) {
    assert(lg[pw[i]] == -1);
    lg[pw[i]] = i;
  }
  if (1) {
    scanf("%d", &n);
    for (int i = 0; i < n; ++i) {
      scanf("%d", &a[i]);
      if (!a[i]) {
        --n;
        --i;
      }
    }
  } else {
    n = 200000;
    for (int i = 0; i < n; ++i) {
      a[i] = 12345;
    }
    cout << 1LL * mul(a[0], a[0]) * n * (n - 1) / 2 << endl;
  }
  for (int i = 0; i < n; ++i) {
    ans -= mul(a[i], a[i]);
    a[i] = lg[a[i]];
    ++A[a[i]];
  }
  for (int i = 0; i < A.size(); ++i) {
    B[i] = A[i] / sq;
    C[i] = A[i] % sq;
  }
  vector<int> c1 = ntt.square(B);
  vector<int> c2 = ntt.product(B, C);
  vector<int> c3 = ntt.square(C);
  for (int i = 0; i < c1.size(); ++i) {
    long long cnt = 1LL * sq * sq * c1[i];
    cnt += 2LL * sq * c2[i];
    cnt += c3[i];
    ans += 1LL * cnt * pw[i % (mod - 1)];
  }
  assert(ans % 2 == 0);
  ans /= 2;
  cout << ans << "\\n";
  return 0;
}


Error: WA

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  for (int i = 0; i < v.size(); ++i) {
    if (i) {
      os << " ";
    }
    os << v[i];
  }
  os << "]";
  return os;
}

const int max_lev = 18, mod_fft = 998244353;

template <int max_lev> struct FFT {
  static const int max_n = 1 << max_lev;

  typedef double real_type;

  template <typename real_type> struct my_complex {
    real_type x, y;

    my_complex(real_type x = 0, real_type y = 0) : x(x), y(y) {}

    void operator+=(const my_complex &o) {
      x += o.x;
      y += o.y;
    }

    void operator-=(const my_complex &o) {
      x -= o.x;
      y -= o.y;
    }

    my_complex operator*(const my_complex &o) const {
      return my_complex(x * o.x - y * o.y, x * o.y + o.x * y);
    }

    void operator/=(real_type k) {
      x /= k;
      y /= k;
    }
  };

  typedef my_complex<real_type> base;

  int rev[max_n];
  real_type si[max_n];
  base buf1[max_n], buf2[max_n], bufpw[max_n];

  FFT() {
    const real_type full = 4 * acosl(0.0);
    for (int i = 0; i < max_n; ++i) {
      si[i] = sinl(full * i / max_n);
      rev[i] = (rev[i / 2] / 2) + ((i & 1) << (max_lev - 1));
    }
  }

  template <bool inv> void fft(base *a, int lg) {
    const int n = 1 << lg;
    for (int i = 0; i < n; ++i) {
      int r = (rev[i] >> (max_lev - lg)) & (n - 1);
      if (i < r) {
        swap(a[i], a[r]);
      }
    }
    int pl = max_n / 2;
    for (int st = 1; st < n; st *= 2) {
      if (inv) {
        for (int i = 0, pos = 0; i < st; ++i, pos += pl) {
          bufpw[i] = {si[(max_n / 4 + pos) & (max_n - 1)], -si[pos]};
        }
      } else {
        for (int i = 0, pos = 0; i < st; ++i, pos += pl) {
          bufpw[i] = {si[(max_n / 4 + pos) & (max_n - 1)], si[pos]};
        }
      }
      for (int i = 0; i < n; i += 2 * st) {
        base *curpw = bufpw;
        for (int j = i; j < i + st; ++j) {
          const base y = *curpw * a[j + st];
          a[j + st] = a[j];
          a[j + st] -= y;
          a[j] += y;
          ++curpw;
        }
      }
      pl /= 2;
    }
  }

  vector<long long> product(const vector<int> &a, const vector<int> &b) {
    int sz = a.size() + b.size() - 1;
    int n = 1, lg = 0;
    while (n < sz) {
      n *= 2;
      ++lg;
    }
    copy(a.begin(), a.end(), buf1);
    copy(b.begin(), b.end(), buf2);
    fill(buf1 + a.size(), buf1 + n, 0);
    fill(buf2 + b.size(), buf2 + n, 0);
    fft<false>(buf1, lg);
    fft<false>(buf2, lg);
    for (int i = 0; i < n; ++i) {
      buf1[i] = buf1[i] * buf2[i];
    }
    fft<true>(buf1, lg);
    vector<long long> ans(sz);
    for (int i = 0; i < sz; ++i) {
      buf1[i] /= n;
      ans[i] = buf1[i].x + 0.5;
    }
    return ans;
  }

  vector<long long> square(const vector<int> &a) {
    int sz = a.size() + a.size() - 1;
    int n = 1, lg = 0;
    while (n < sz) {
      n *= 2;
      ++lg;
    }
    copy(a.begin(), a.end(), buf1);
    fill(buf1 + a.size(), buf1 + n, 0);
    fft<false>(buf1, lg);
    for (int i = 0; i < n; ++i) {
      buf1[i] = buf1[i] * buf1[i];
    }
    fft<true>(buf1, lg);
    vector<long long> ans(sz);
    for (int i = 0; i < sz; ++i) {
      buf1[i] /= n;
      ans[i] = buf1[i].x + 0.5;
    }
    return ans;
  }
};

FFT<max_lev> fft;

const int max_n = 200222, inf = 1000111222;
const int mod = 200003;
const int g = 2;

int mul(int x, int y) { return (1LL * x * y) % mod; }

int n, a[max_n];
int pw[max_n], lg[max_n];
long long ans;

vector<int> A(mod), B(mod), C(mod);

int main() {
  // freopen("input.txt", "r", stdin);
  // freopen("output.txt", "w", stdout);
  pw[0] = 1;
  for (int i = 1; i + 1 < mod; ++i) {
    pw[i] = mul(pw[i - 1], g);
  }
  memset(lg, -1, sizeof(lg));
  for (int i = 0; i + 1 < mod; ++i) {
    assert(lg[pw[i]] == -1);
    lg[pw[i]] = i;
  }
  if (1) {
    scanf("%d", &n);
    for (int i = 0; i < n; ++i) {
      scanf("%d", &a[i]);
      if (!a[i]) {
        --n;
        --i;
      }
    }
  } else {
    n = 200000;
    for (int i = 0; i < n; ++i) {
      a[i] = 12345;
    }
    cout << 1LL * mul(a[0], a[0]) * n * (n - 1) / 2 << endl;
  }
  for (int i = 0; i < n; ++i) {
    ans -= mul(a[i], a[i]);
    a[i] = lg[a[i]];
    ++A[a[i]];
  }
  vector<long long> c1 = fft.square(A);
  for (int i = 0; i < c1.size(); ++i) {
    long long cnt = c1[i];
    ans += 1LL * cnt * pw[i % (mod - 1)];
  }
  assert(ans % 2 == 0);
  ans /= 2;
  cout << ans << "\\n";
  return 0;
}