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

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
using namespace std;
const double pi = acos(-1);
int n, val, cnt, sz = 4194304, p[4194333], d[4194333], c[4194333];
complex<double> a[4194333], b[4194333], y[4194333];
void dft(complex<double> *x, bool inv) {
  for (int i = sz >> 1; i >= 1; i >>= 1) {
    complex<double> bs(cos(2.0 * pi * i / sz),
                       (inv ? -1.0 : 1.0) * sin(2.0 * pi * i / sz)),
        w = 1;
    for (int j = 0; 2 * j < sz; j += i) {
      for (int k = 0; k < i; k++) {
        y[j + k] = x[2 * j + k] + w * x[2 * j + k + i];
        y[j + k + sz / 2] = x[2 * j + k] - w * x[2 * j + k + i];
      }
      w *= bs;
    }
    for (int j = 0; j < sz; j++)
      x[j] = y[j];
  }
  if (inv)
    for (int i = 0; i < sz; i++)
      x[i] /= sz;
}
int main() {
  scanf("%d", &n);
  int s = 2 * n * n;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      scanf("%d", &val);
      a[2 * i * n + j] = val;
      b[s - (2 * i * n + j) - 1] = val;
      cnt += val;
    }
  }
  dft(a, false);
  dft(b, false);
  for (int i = 0; i < sz; i++)
    a[i] *= b[i];
  dft(a, true);
  for (int i = 0; i < sz; i++)
    c[i] = a[i].real() + 0.5;
  double sum = 0.0;
  for (int i = 0; i < 2 * s; i++) {
    int ri = i - s + 1;
    int rx = ((ri + s) % (2 * n) < n ? (ri + s) % (2 * n)
                                     : 2 * n - (ri + s) % (2 * n));
    int ry = ((ri + s + n) / (2 * n) >= n ? (ri + s + n) / (2 * n) - n
                                          : n - (ri + s + n) / (2 * n));
    d[i] = rx * rx + ry * ry;
    sum += sqrt(d[i]) * c[i];
    p[i] = i;
  }
  sort(p, p + 2 * s, [](int i, int j) { return d[i] < d[j]; });
  printf("%.12lf\\n", sum / cnt / (cnt - 1));
  long long cur = 0;
  int counter = 10000;
  for (int i = 1; i <= 2 * s; i++) {
    cur += c[p[i - 1]];
    if (i == 2 * s || d[p[i]] != d[p[i - 1]]) {
      long long realcur = (cur - (d[p[i - 1]] == 0 ? cnt : 0)) / 2;
      if (realcur > 0 && counter-- > 0)
        printf("%d %lld\\n", d[p[i - 1]], realcur);
      cur = 0;
    }
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#pragma warning(disable : 4996)
using namespace std;
// Naive TLE Code (it will get TLE)
int n, val, cnt, sz = 2097152, a[2097177], b[2097177], c[4194333], p[4194333],
                 d[4194333];
int main() {
  scanf("%d", &n);
  int s = 2 * n * n;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      scanf("%d", &val);
      a[2 * i * n + j] = val;
      b[s - (2 * i * n + j) - 1] = val;
      cnt += val;
    }
  }
  for (int i = 0; i < s; i++) {
    for (int j = 0; j <= s; j++) {
      c[i + j] += a[i] * b[j];
    }
  }
  double sum = 0.0;
  for (int i = 0; i < 2 * s; i++) {
    int ri = i - s + 1;
    int rx = ((ri + s) % (2 * n) < n ? (ri + s) % (2 * n)
                                     : 2 * n - (ri + s) % (2 * n));
    int ry = ((ri + s + n) / (2 * n) >= n ? (ri + s + n) / (2 * n) - n
                                          : n - (ri + s + n) / (2 * n));
    d[i] = rx * rx + ry * ry;
    sum += sqrt(d[i]) * c[i];
    p[i] = i;
  }
  sort(p, p + 2 * s, [](int i, int j) { return d[i] < d[j]; });
  printf("%.12lf\\n", sum / cnt / (cnt - 1));
  long long cur = 0;
  int counter = 10000;
  for (int i = 1; i <= 2 * s; i++) {
    cur += c[p[i - 1]];
    if (i == 2 * s || d[p[i]] != d[p[i - 1]]) {
      long long realcur = (cur - (d[p[i - 1]] == 0 ? cnt : 0)) / 2;
      if (realcur > 0 && counter-- > 0)
        printf("%d %lld\\n", d[p[i - 1]], realcur);
      cur = 0;
    }
  }
  return 0;
}