AOJ 2560: Point Distance

二次元になっているものを一列にしてFFTをする。後は重複回数をよく考えつつうまくFFTの表から答えを出す。

実行時間がかなりかかった。なぜだろう。

#include<stdio.h>
#include<algorithm>
#include<complex>
#include<vector>
#include<math.h>
#include<map>
using namespace std;
const double PI = 4.0*atan(1.0);
typedef complex<double> Complex;
const Complex I(0, 1);
void fft(int n, double theta, Complex a[]) {
  for (int m = n; m >= 2; m >>= 1) {
    int mh = m >> 1;
    for (int i = 0; i < mh; i++) {
      Complex w = exp(i*theta*I);
      for (int j = i; j < n; j += m) {
        int k = j + mh;
        Complex x = a[j] - a[k];
        a[j] += a[k];
        a[k] = w * x;
      }
    }
    theta *= 2;
  }
  int i = 0;
  for (int j = 1; j < n - 1; j++) {
    for (int k = n >> 1; k > (i ^= k); k >>= 1);
    if (j < i) swap(a[i], a[j]);
  }
}
int b[1025][1025];
Complex p[1<<22];
Complex q[1<<22];
int main(){
	int a;
	scanf("%d",&a);
	int cnt=0;
	for(int i=0;i<a;i++){
		for(int j=0;j<a;j++)scanf("%d",&b[i][j]);
	}
	int n=1;
	while(n<a*2)n*=2;
	for(int i=0;i<a;i++)for(int j=0;j<a;j++){
		cnt+=b[i][j];
		p[i*n+j]=q[(n-1-i)*n+(n-1-j)]=Complex(b[i][j],0);
	}
	fft(n*n,2*PI/n/n,p);
	fft(n*n,2*PI/n/n,q);
	
	for(int i=0;i<n*n;i++)p[i]=p[i]*q[i];
	fft(n*n,-2*PI/n/n,p);
	for(int i=0;i<n*n;i++)p[i]/=n*n;
	/*for(int i=0;i<n;i++){
		for(int j=0;j<n;j++)printf("%.0f ",p[i*n+j].real());
		printf("\n");
	}*/
	double ret=0;
	long long sz=0;
	map<int,long long>m;
	for(int i=0;i<n*n;i++){
			int pi=i/n;
			int pj=i%n;
			if(pj>=n/2)continue;
			if(pj==0&&pi>=n/2)continue;
			long long v=(long long)(p[n*n-1-i].real()+0.5);
			if(i==0){v-=cnt;v/=2;}
			if(!v)continue;
			
			int I=min(pi,n-pi);
			int J=pj;
			int t=I*I+J*J;
			if(m.count(t)){
				m[t]+=v;
			}else{
				m[t]=v;
			}
			sz+=v;
			ret+=sqrt(t)*v;
	}
	printf("%.12f\n",ret/sz);
	int at=0;
	for(map<int,long long>::iterator it=m.begin();it!=m.end();it++){
		if(at==10000)break;
		at++;
		printf("%d %lld\n",(*it).first,(*it).second);
	}
}