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

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)
#define all(c) c.begin(), c.end()
#define pb push_back
#define fs first
#define sc second
#define show(x) cout << #x << " = " << x << endl
#define chmin(x, y) x = min(x, y)
#define chmax(x, y) x = max(x, y)
using namespace std;

typedef long long ll;
typedef pair<ll, ll> P; // val,mod;
ll inv(ll a, ll m) {
  ll b = m, u = 1, v = 0;
  while (b) {
    ll t = a / b;
    swap(a -= t * b, b);
    swap(u -= t * v, v);
  }
  u %= m;
  if (u < 0)
    u += m;
  return u;
}

struct NTT {
  const int mod, g;
  NTT(int mod, int g) : mod(mod), g(g) {}
  ll ex(ll x, ll p) {
    ll a = 1;
    while (p) {
      if (p % 2)
        a = a * x % mod;
      x = x * x % mod;
      p /= 2;
    }
    return a;
  }
  ll inv(ll a) {
    ll b = mod, u = 1, v = 0;
    while (b) {
      ll t = a / b;
      swap(a -= t * b, b);
      swap(u -= t * v, v);
    }
    u %= mod;
    if (u < 0)
      u += mod;
    return u;
  }
  void ntt(vector<ll> &a, ll h) {
    int N = a.size();
    for (int m = N; m >= 2; m /= 2) {
      int mh = m / 2;
      ll w = 1;
      rep(i, mh) {
        for (int j = i; j < N; j += m) {
          int k = j + mh;
          ll x = (a[j] - a[k]);
          a[j] += a[k];
          if (a[j] >= mod)
            a[j] -= mod;
          a[k] = w * x % mod;
          if (a[k] < 0)
            a[k] += mod;
        }
        w = w * h % mod;
      }
      h = h * h % mod;
    }
    int i = 0;
    rep1(j, N - 2) {
      for (int k = N / 2; k > (i ^= k); k /= 2)
        ;
      if (j < i)
        swap(a[i], a[j]);
    }
  }
  vector<ll> conv(vector<ll> a,
                  vector<ll> b) { // (mod-1)%N==0??§????????¨??????!!
    for (auto &x : a)
      x %= mod;
    for (auto &x : b)
      x %= mod;
    int n = a.size() + b.size();
    int N = 1;
    while (N < n)
      N *= 2;
    a.resize(N), b.resize(N);
    assert((mod - 1) % N == 0);
    ll h = ex(g, (mod - 1) / N);
    ntt(a, h), ntt(b, h);
    rep(i, N)(a[i] *= b[i]) %= mod;
    h = inv(h);
    ntt(a, h);
    h = inv(N);
    rep(i, N)(a[i] *= h) %= mod;
    return a;
  }
};

int main() {
  int N;
  cin >> N;
  vector<ll> A(N * N * 2), B;
  rep(i, N) rep(j, N) { cin >> A[i * 2 * N + j]; }
  B = A;
  reverse(all(B));
  auto C = NTT(1224736769, 3).conv(A, B);
  vector<ll> mp(1 << 21);
  rep(i, N) for (int j = -N + 1; j <= N - 1; j++) {
    if (i == 0 && j <= 0)
      continue;
    mp[i * i + j * j] += C[2 * N * N - 1 - i * 2 * N + j];
  }
  rep(i, N) rep(j, N) {
    int x = A[i * 2 * N + j];
    mp[0] += x * (x - 1) / 2;
  }
  long long sum = 0;
  double ave = 0;
  rep(i, N * N * 2) {
    if (mp[i] == 0)
      continue;
    ave += sqrt(i) * mp[i];
    sum += mp[i];
  }
  printf("%.12f\\n", ave / sum);
  int c = 0;
  rep(i, N * N * 2) {
    if (mp[i] == 0)
      continue;
    printf("%d %lld\\n", i, mp[i]);
    c++;
    if (c == 10000)
      break;
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)
#define all(c) c.begin(), c.end()
#define pb push_back
#define fs first
#define sc second
#define show(x) cout << #x << " = " << x << endl
#define chmin(x, y) x = min(x, y)
#define chmax(x, y) x = max(x, y)
using namespace std;

typedef long long ll;
typedef pair<ll, ll> P; // val,mod;
ll inv(ll a, ll m) {
  ll b = m, u = 1, v = 0;
  while (b) {
    ll t = a / b;
    swap(a -= t * b, b);
    swap(u -= t * v, v);
  }
  u %= m;
  if (u < 0)
    u += m;
  return u;
}

struct NTT {
  const int mod, g;
  NTT(int mod, int g) : mod(mod), g(g) {}
  ll ex(ll x, ll p) {
    ll a = 1;
    while (p) {
      if (p % 2)
        a = a * x % mod;
      x = x * x % mod;
      p /= 2;
    }
    return a;
  }
  ll inv(ll a) {
    ll b = mod, u = 1, v = 0;
    while (b) {
      ll t = a / b;
      swap(a -= t * b, b);
      swap(u -= t * v, v);
    }
    u %= mod;
    if (u < 0)
      u += mod;
    return u;
  }
  void ntt(vector<ll> &a, ll h) {
    int N = a.size();
    for (int m = N; m >= 2; m /= 2) {
      int mh = m / 2;
      ll w = 1;
      rep(i, mh) {
        for (int j = i; j < N; j += m) {
          int k = j + mh;
          ll x = (a[j] - a[k]);
          a[j] += a[k];
          if (a[j] >= mod)
            a[j] -= mod;
          a[k] = w * x % mod;
          if (a[k] < 0)
            a[k] += mod;
        }
        w = w * h % mod;
      }
      h = h * h % mod;
    }
    int i = 0;
    rep1(j, N - 2) {
      for (int k = N / 2; k > (i ^= k); k /= 2)
        ;
      if (j < i)
        swap(a[i], a[j]);
    }
  }
  vector<ll> conv(vector<ll> a,
                  vector<ll> b) { // (mod-1)%N==0??§????????¨??????!!
    for (auto &x : a)
      x %= mod;
    for (auto &x : b)
      x %= mod;
    int n = a.size() + b.size();
    int N = 1;
    while (N < n)
      N *= 2;
    a.resize(N), b.resize(N);
    assert((mod - 1) % N == 0);
    ll h = ex(g, (mod - 1) / N);
    ntt(a, h), ntt(b, h);
    rep(i, N)(a[i] *= b[i]) %= mod;
    h = inv(h);
    ntt(a, h);
    h = inv(N);
    rep(i, N)(a[i] *= h) %= mod;
    return a;
  }
};

vector<ll> ms = {167772161, 469762049}, gs = {3, 3};
vector<ll> bigconv(vector<ll> a, vector<ll> b) {
  vector<ll> vs[2];
  rep(i, 2) vs[i] = NTT(ms[i], gs[i]).conv(a, b);
  int N = vs[0].size();
  vector<ll> ret(N);
  ll iv = inv(ms[0], ms[1]);
  rep(i, N) {
    ll tmp = (vs[1][i] - vs[0][i]) * iv % ms[1];
    if (tmp < 0)
      tmp += ms[1];
    ret[i] = vs[0][i] + ms[0] * tmp;
  }
  return ret;
}

int main() {
  int N;
  cin >> N;
  vector<ll> A(N * N * 2), B;
  rep(i, N) rep(j, N) { cin >> A[i * 2 * N + j]; }
  B = A;
  reverse(all(B));
  auto C = bigconv(A, B);
  vector<ll> mp(1 << 21);
  rep(i, N) for (int j = -N + 1; j <= N - 1; j++) {
    if (i == 0 && j <= 0)
      continue;
    mp[i * i + j * j] += C[2 * N * N - 1 - i * 2 * N + j];
  }
  rep(i, N) rep(j, N) {
    int x = A[i * 2 * N + j];
    mp[0] += x * (x - 1) / 2;
  }
  long long sum = 0;
  double ave = 0;
  rep(i, N * N * 2) {
    if (mp[i] == 0)
      continue;
    ave += sqrt(i) * mp[i];
    sum += mp[i];
  }
  printf("%.12f\\n", ave / sum);
  int c = 0;
  rep(i, N * N * 2) {
    if (mp[i] == 0)
      continue;
    printf("%d %lld\\n", i, mp[i]);
    c++;
    if (c == 10000)
      break;
  }
}