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

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <map>
#include <vector>

using namespace std;

typedef double Real;
typedef vector<Real> VR;
typedef complex<Real> Cpx;
typedef vector<Cpx> VC;

const Real PI = acos(-1.0);

struct FFT {
  // FFT with trigonometry
  Cpx pool[262144];
  int c;
  VC fft(VC &A, Cpx w) {
    if (A.size() == 1)
      return A;
    VC Aeven, Aodd;
    for (int i = 0; i < A.size(); i++) {
      if (i % 2 == 0)
        Aeven.push_back(A[i]);
      else
        Aodd.push_back(A[i]);
    }
    VC Feven = fft(Aeven, w * w);
    VC Fodd = fft(Aodd, w * w);
    VC F(A.size());
    Cpx x = 1;
    for (int i = 0; i < A.size() / 2; i++) {
      F[i] = Feven[i] + Fodd[i] * x;
      F[i + A.size() / 2] = Feven[i] - Fodd[i] * x;
      x *= w;
    }
    return F;
  }
  int fft(int id, int ln, Cpx w) {
    if (ln == 1)
      return id;
    int st = c;
    // even
    for (int i = 0; i < ln; i += 2) {
      pool[c++] = pool[id + i];
    }
    int Feven = fft(st, ln / 2, w * w);
    // odd
    st = c;
    for (int i = 1; i < ln; i += 2) {
      pool[c++] = pool[id + i];
    }
    int Fodd = fft(st, ln / 2, w * w);
    // res
    int res = c;
    Cpx x = 1;
    for (int i = 0; i < ln / 2; i++) {
      pool[c + i] = pool[Feven + i] + pool[Fodd + i] * x;
      pool[c + i + ln / 2] = pool[Feven + i] - pool[Fodd + i] * x;
      x *= w;
    }
    c += ln;
    return res;
  }
  /*	VC fft2(VC A,int n=-1){
                  if(n==-1) n=A.size();
                  int N=1;
                  while(N<n) N*=2;
                  for(int i=A.size();i<N;i++) A.push_back(0);
                  VC res=fft(A,Cpx(cos(-PI*2/N),sin(-PI*2/N)));
                  return res;
          }*/
  VC fft2(VC &A, int n = -1) {
    c = 0;
    if (n == -1)
      n = A.size();
    int N = 1;
    while (N < n)
      N *= 2;
    for (int i = 0; i < A.size(); i++)
      pool[c++] = A[i];
    for (int i = A.size(); i < N; i++)
      pool[c++] = 0;
    int id = fft(0, N, Cpx(cos(-PI * 2 / N), sin(-PI * 2 / N)));
    VC res;
    for (int i = 0; i < N; i++)
      res.push_back(pool[id + i]);
    return res;
  }
  /*	VC revFFT2(VC A,int n=-1){
                  if(n==-1) n=A.size();
                  int N=1;
                  while(N<n) N*=2;
                  for(int i=A.size();i<N;i++) A.push_back(0);
                  VC res=fft(A,Cpx(cos(PI*2/N),sin(PI*2/N)));
                  for(int i=0;i<res.size();i++) res[i]/=res.size();
                  return res;
          }*/
  VC revFFT2(VC &A, int n = -1) {
    c = 0;
    if (n == -1)
      n = A.size();
    int N = 1;
    while (N < n)
      N *= 2;
    for (int i = 0; i < A.size(); i++)
      pool[c++] = A[i];
    for (int i = A.size(); i < N; i++)
      pool[c++] = 0;
    int id = fft(0, N, Cpx(cos(PI * 2 / N), sin(PI * 2 / N)));
    VC res;
    for (int i = 0; i < N; i++)
      res.push_back(pool[id + i] / (Real)N);
    return res;
  }
};

FFT fft;

Cpx C[2][2048][2048];
int N;

Cpx res[2048][2048];

long long ans[1024 * 1024 * 2];

int main() {
  Real sum = 0;
  scanf("%d", &N);
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      Real x;
      scanf("%lf", &x);
      sum += x;
      C[0][i][j] = Cpx(x, 0);
    }
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      C[1][N - 1 - i][N - 1 - j] = C[0][i][j];
    }
  int N_ = N;
  N = 1;
  while (N < N_ * 2)
    N *= 2;
  for (int i = 0; i < 2; i++) {
    for (int j = 0; j < N; j++) {
      VC A;
      for (int k = 0; k < N; k++)
        A.push_back(C[i][j][k]);
      VC FA = fft.fft2(A);
      for (int k = 0; k < N; k++)
        C[i][j][k] = FA[k];
    }
  }
  for (int i = 0; i < 2; i++) {
    for (int k = 0; k < N; k++) {
      VC A;
      for (int j = 0; j < N; j++)
        A.push_back(C[i][j][k]);
      VC FA = fft.fft2(A);
      for (int j = 0; j < N; j++)
        C[i][j][k] = FA[j];
    }
  }
  /*for(int i=0;i<2;i++){
          for(int j=0;j<N;j++){
                  for(int k=0;k<N;k++){
                          printf("(%.2f %.2f)
  ",C[i][j][k].real(),C[i][j][k].imag());
                  }
                  printf("\\n");
          }
  }*/
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      res[i][j] = C[0][i][j] * C[1][i][j];
    }
  for (int j = 0; j < N; j++) {
    VC A;
    for (int i = 0; i < N; i++)
      A.push_back(res[i][j]);
    VC FA = fft.revFFT2(A);
    for (int i = 0; i < N; i++)
      res[i][j] = FA[i];
  }
  for (int i = 0; i < N; i++) {
    VC A;
    for (int j = 0; j < N; j++)
      A.push_back(res[i][j]);
    VC FA = fft.revFFT2(A);
    for (int j = 0; j < N; j++)
      res[i][j] = FA[j];
  }
  Real num = 0, den = 0;
  for (int i = 0; i < N_ * 2 - 1; i++)
    for (int j = 0; j < N_ * 2 - 1; j++) {
      int di = N_ - 1 - i;
      int dj = N_ - 1 - j;
      int cnt = (int)(res[i][j] + 0.1).real();
      num += sqrt(di * di + dj * dj) * cnt;
      den += cnt;
      ans[di * di + dj * dj] += cnt;
    }
  den -= sum;
  printf("%.9f\\n", num / den);
  int c = 0;
  ans[0] -= (int)(sum + 0.1);
  for (int i = 0; i < 1024 * 1024 * 2; i++) {
    if (ans[i] == 0)
      continue;
    ans[i] /= 2;
    printf("%d %lld\\n", i, ans[i]);
    c++;
    if (c == 10000)
      break;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <map>
#include <vector>

using namespace std;

typedef double Real;
typedef vector<Real> VR;
typedef complex<Real> Cpx;
typedef vector<Cpx> VC;

const Real PI = acos(-1.0);

struct FFT {
  // FFT with trigonometry
  VC fft(VC &A, Cpx w) {
    if (A.size() == 1)
      return A;
    VC Aeven, Aodd;
    for (int i = 0; i < A.size(); i++) {
      if (i % 2 == 0)
        Aeven.push_back(A[i]);
      else
        Aodd.push_back(A[i]);
    }
    VC Feven = fft(Aeven, w * w);
    VC Fodd = fft(Aodd, w * w);
    VC F(A.size());
    Cpx x = 1;
    for (int i = 0; i < A.size() / 2; i++) {
      F[i] = Feven[i] + Fodd[i] * x;
      F[i + A.size() / 2] = Feven[i] - Fodd[i] * x;
      x *= w;
    }
    return F;
  }
  VC fft2(VC A, int n = -1) {
    if (n == -1)
      n = A.size();
    int N = 1;
    while (N < n)
      N *= 2;
    for (int i = A.size(); i < N; i++)
      A.push_back(0);
    VC res = fft(A, Cpx(cos(-PI * 2 / N), sin(-PI * 2 / N)));
    return res;
  }
  VC revFFT2(VC A, int n = -1) {
    if (n == -1)
      n = A.size();
    int N = 1;
    while (N < n)
      N *= 2;
    for (int i = A.size(); i < N; i++)
      A.push_back(0);
    VC res = fft(A, Cpx(cos(PI * 2 / N), sin(PI * 2 / N)));
    for (int i = 0; i < res.size(); i++)
      res[i] /= res.size();
    return res;
  }
};

FFT fft;

Cpx C[2][2048][2048];
int N;

Cpx res[2048][2048];

long long ans[1024 * 1024 * 2];

int main() {
  Real sum = 0;
  scanf("%d", &N);
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      Real x;
      scanf("%lf", &x);
      sum += x;
      C[0][i][j] = Cpx(x, 0);
    }
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      C[1][N - 1 - i][N - 1 - j] = C[0][i][j];
    }
  int N_ = N;
  N = 1;
  while (N < N_ * 2)
    N *= 2;
  for (int i = 0; i < 2; i++) {
    for (int j = 0; j < N; j++) {
      VC A;
      for (int k = 0; k < N; k++)
        A.push_back(C[i][j][k]);
      VC FA = fft.fft2(A);
      for (int k = 0; k < N; k++)
        C[i][j][k] = FA[k];
    }
  }
  for (int i = 0; i < 2; i++) {
    for (int k = 0; k < N; k++) {
      VC A;
      for (int j = 0; j < N; j++)
        A.push_back(C[i][j][k]);
      VC FA = fft.fft2(A);
      for (int j = 0; j < N; j++)
        C[i][j][k] = FA[j];
    }
  }
  /*for(int i=0;i<2;i++){
          for(int j=0;j<N;j++){
                  for(int k=0;k<N;k++){
                          printf("(%.2f %.2f)
  ",C[i][j][k].real(),C[i][j][k].imag());
                  }
                  printf("\\n");
          }
  }*/
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
      res[i][j] = C[0][i][j] * C[1][i][j];
    }
  for (int j = 0; j < N; j++) {
    VC A;
    for (int i = 0; i < N; i++)
      A.push_back(res[i][j]);
    VC FA = fft.revFFT2(A);
    for (int i = 0; i < N; i++)
      res[i][j] = FA[i];
  }
  for (int i = 0; i < N; i++) {
    VC A;
    for (int j = 0; j < N; j++)
      A.push_back(res[i][j]);
    VC FA = fft.revFFT2(A);
    for (int j = 0; j < N; j++)
      res[i][j] = FA[j];
  }
  Real num = 0, den = 0;
  for (int i = 0; i < N_ * 2 - 1; i++)
    for (int j = 0; j < N_ * 2 - 1; j++) {
      int di = N_ - 1 - i;
      int dj = N_ - 1 - j;
      int cnt = (int)(res[i][j] + 0.1).real();
      num += sqrt(di * di + dj * dj) * cnt;
      den += cnt;
      ans[di * di + dj * dj] += cnt;
    }
  den -= sum;
  printf("%.9f\\n", num / den);
  int c = 0;
  ans[0] -= (int)(sum + 0.1);
  for (int i = 0; i < 1024 * 1024 * 2; i++) {
    if (ans[i] == 0)
      continue;
    ans[i] /= 2;
    printf("%d %lld\\n", i, ans[i]);
    c++;
    if (c == 10000)
      break;
  }
  return 0;
}