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

#include <cassert>
#include <ccomplex>
#include <iostream>
#include <vector>

using namespace std;

const double PI = acos(-1);

typedef complex<double> Comp;

bool is_pow_of2(int n) {
  int a = 1;
  while (n >= a) {
    if (n == a) {
      return true;
    }
    a <<= 1;
  }
  return false;
}

vector<Comp> fft(vector<Comp> v) {
  int n = v.size();
  assert(is_pow_of2(n));

  for (int j = 1, i = 0; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1)
      ;
    if (j < i)
      swap(v[i], v[j]);
  }

  for (int m = 2; m <= n; m <<= 1) {
    double deg = (-1) * 2 * PI / m;
    Comp r(cos(deg), sin(deg));
    for (int i = 0; i < n; i += m) {
      Comp w(1, 0);
      for (int j = i, k = i + m / 2; k < i + m; j++, k++) {
        Comp t1 = v[j], t2 = w * v[k];
        v[j] = t1 + t2, v[k] = t1 - t2;
        w *= r;
      }
    }
  }
  return v;
}

vector<Comp> ifft(vector<Comp> v) {
  int n = v.size();
  assert(is_pow_of2(n));
  for (int j = 1, i = 0; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1)
      ;
    if (j < i)
      swap(v[i], v[j]);
  }
  for (int m = 2; m <= n; m <<= 1) {
    double deg = 2 * PI / m;
    Comp r(cos(deg), sin(deg));
    for (int i = 0; i < n; i += m) {
      Comp w(1, 0);
      for (int j = i, k = i + m / 2; k < i + m; j++, k++) {
        Comp t1 = v[j], t2 = w * v[k];
        v[j] = t1 + t2, v[k] = t1 - t2;
        w *= r;
      }
    }
  }
  for (int i = 0; i < n; ++i)
    v[i] *= 1.0 / n;
  return v;
}

vector<Comp> convolution(vector<Comp> P, vector<Comp> Q) {
  P = fft(P), Q = fft(Q);
  for (int i = 0; i < P.size(); i++)
    P[i] *= Q[i];
  return ifft(P);
}

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0, i##_len = n; i < i##_len; i++)
#define all(x) begin(x), end(x)
#define lim(x, r, l) (r <= x && x < l)

#define int long long

signed main() {
  int n;
  cin >> n;
  vector<vector<int>> c(n, vector<int>(n, 0));
  int zero = 0;
  rep(i, n) rep(j, n) {
    int cc;
    cin >> cc;
    c[i][j] = cc;
    if (cc >= 2)
      zero += (cc - 1) * (cc) / 2;
  }
  vector<Comp> a(1 << 22), b(1 << 22);

  rep(i, n) {
    rep(j, n) {
      a[2 * n * i + j] = c[i][j];
      b[2 * n * i + j] = c[n - i - 1][n - j - 1];
    }
  }

  vector<Comp> res = convolution(a, b);

  double sum = 0;
  int num = 0;
  map<int, int> ans;

  rep(i, res.size()) {
    int px = n - i / (2 * n) - 1;
    int py = n - i % (2 * n) - 1;
    int s = real(res[i]) + 0.5;
    int len = px * px + py * py;
    // cout << s << endl;
    if (s == 0) {
      continue;
    }
    if (len == 0) {
      continue;
    }
    if (ans.find(len) != ans.end()) {
      ans[len] += s;
    } else {
      ans[len] = s;
    }
    sum += sqrt(len) * s;
    num += s;
  }

  cout << setprecision(12) << sum / 2 / (num / 2 + zero) << endl;
  int ii = 0;
  if (zero > 0) {
    cout << "0 " << zero << endl;
    ii++;
  }
  for (auto aa : ans) {
    if (ii >= 10000)
      break;
    int l, s;
    tie(l, s) = aa;
    cout << l << " " << s / 2 << endl;
    ii++;
  }
}

Error: MLE

Fixed submission:

#include <cassert>
#include <ccomplex>
#include <iostream>
#include <vector>

using namespace std;

const double PI = acos(-1);

typedef complex<double> Comp;

bool is_pow_of2(int n) {
  int a = 1;
  while (n >= a) {
    if (n == a) {
      return true;
    }
    a <<= 1;
  }
  return false;
}

vector<Comp> fft(vector<Comp> v) {
  int n = v.size();
  assert(is_pow_of2(n));

  for (int j = 1, i = 0; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1)
      ;
    if (j < i)
      swap(v[i], v[j]);
  }

  for (int m = 2; m <= n; m <<= 1) {
    double deg = (-1) * 2 * PI / m;
    Comp r(cos(deg), sin(deg));
    for (int i = 0; i < n; i += m) {
      Comp w(1, 0);
      for (int j = i, k = i + m / 2; k < i + m; j++, k++) {
        Comp t1 = v[j], t2 = w * v[k];
        v[j] = t1 + t2, v[k] = t1 - t2;
        w *= r;
      }
    }
  }
  return v;
}

vector<Comp> ifft(vector<Comp> v) {
  int n = v.size();
  assert(is_pow_of2(n));
  for (int j = 1, i = 0; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1)
      ;
    if (j < i)
      swap(v[i], v[j]);
  }
  for (int m = 2; m <= n; m <<= 1) {
    double deg = 2 * PI / m;
    Comp r(cos(deg), sin(deg));
    for (int i = 0; i < n; i += m) {
      Comp w(1, 0);
      for (int j = i, k = i + m / 2; k < i + m; j++, k++) {
        Comp t1 = v[j], t2 = w * v[k];
        v[j] = t1 + t2, v[k] = t1 - t2;
        w *= r;
      }
    }
  }
  for (int i = 0; i < n; ++i)
    v[i] *= 1.0 / n;
  return v;
}

vector<Comp> convolution(vector<Comp> P, vector<Comp> Q) {
  P = fft(P), Q = fft(Q);
  for (int i = 0; i < P.size(); i++)
    P[i] *= Q[i];
  return ifft(P);
}

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0, i##_len = n; i < i##_len; i++)
#define all(x) begin(x), end(x)
#define lim(x, r, l) (r <= x && x < l)

int main() {
  int n;
  cin >> n;
  vector<vector<int>> c(n, vector<int>(n, 0));
  int zero = 0;
  rep(i, n) rep(j, n) {
    int cc;
    cin >> cc;
    c[i][j] = cc;
    if (cc >= 2)
      zero += (cc - 1) * (cc) / 2;
  }
  vector<Comp> a(1 << 23), b(1 << 23);

  rep(i, n) {
    rep(j, n) {
      a[2 * n * i + j] = c[i][j];
      b[2 * n * (i + 1) + (j + 1)] = c[n - i - 1][n - j - 1];
    }
  }

  vector<Comp> res = convolution(a, b);

  double sum = 0;
  int num = 0;
  map<int, int> ans;

  rep(i, (2 * n + 1) * n) {
    int px = n - i / (2 * n);
    int py = n - i % (2 * n);
    int s = real(res[i]) + 0.5;
    int len = px * px + py * py;
    // cout << s << endl;
    if (s == 0) {
      continue;
    }
    if (len == 0) {
      continue;
    }
    if (ans.find(len) != ans.end()) {
      ans[len] += s;
    } else {
      ans[len] = s;
    }
    sum += sqrt(len) * s;
    num += s;
  }

  cout << setprecision(12) << sum / (num + zero) << endl;
  int ii = 0;
  if (zero > 0) {
    cout << "0 " << zero << endl;
    ii++;
  }
  for (auto aa : ans) {
    if (ii >= 10000)
      break;
    int l, s;
    tie(l, s) = aa;
    cout << l << " " << s << endl;
    ii++;
  }
}