2次元平面の各格子点上に9個以下の点がある。(本質だけに絞ると)距離が√0、√1、√2・・・になる点のペアがそれぞれ何個ずつあるかを答えよ。座標のサイズは縦横1024以下。
格子点(x,y)にある点の個数をA[x][y]とすると、ようするにA[0][0]A[i][j]+A[0][1]A[i][j+1]+・・・+A[1][0]A[i+1][j]+・・・を各i,jに対して求めたいんだけど、これは見た目としてFFTが使えそうな気がしてくる。
実際、Aの縦横を逆にしたものをBとすると上の式はA[0][0]B[n-1-i][n-1-j]+A[0][1]B[n-1-i][n-2-j]+・・・+A[1][0]B[n-2-i][n-1-j]+・・・となるので畳み込み演算の形になる。
よって、
1.A,Bを縦方向にFFTする
2.A,Bを横方向にFFTする
3.各i,jに対してA[i][j]*=B[i][j]を行う
4.Aを横方向に逆FFTする
5.Aを縦方向に逆FFTする
ってすれば求めたいものが求まる。
//JAG春コンテスト13F PointDistance(AOJ2560)
#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define pf push_front
typedef long long lint;
typedef complex<double> P;
#define mp make_pair
#define fi first
#define se second
typedef pair<lint,lint> pint;
#define All(s) s.begin(),s.end()
#define rAll(s) s.rbegin(),s.rend()
#define REP(i,a,b) for(int i=a;i<b;i++)
#define rep(i,n) REP(i,0,n)
#define N 2048
const double pi = 4.0*atan(1.0);
const P I(0, 1);
vector<P> fft(int n, double theta, vector<P> a) {
for (int m = n; m >= 2; m >>= 1) {
int mh = m >> 1;
for (int i = 0; i < mh; i++) {
P w = exp(i*theta*I);
for (int j = i; j < n; j += m) {
int k = j + mh;
P 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]);
}
return a;
}
int ma[1030][1030];
map<lint,lint> me;
vector<pint> out;
vector<vector<P> > a(N,vector<P>(N,P(0,0))),b(N,vector<P>(N,P(0,0)));
int main()
{
int n;
cin>>n;
rep(i,n) rep(j,n) cin>>ma[i][j];
rep(i,n) rep(j,n){
a[i][j].real()=ma[i][j];
b[n-i-1][n-j-1].real()=ma[i][j];
}
rep(i,N) a[i]=fft(N,2*pi/N,a[i]),b[i]=fft(N,2*pi/N,b[i]);
rep(i,N){
vector<P> x,y;
rep(j,N) x.pb(a[j][i]),y.pb(b[j][i]);
x=fft(N,2*pi/N,x);y=fft(N,2*pi/N,y);
rep(j,N) a[j][i]=x[j],b[j][i]=y[j];
}
rep(i,N) rep(j,N) a[i][j]*=b[i][j];
rep(i,N){
vector<P> x;
rep(j,N) x.pb(a[j][i]);
x=fft(N,-2*pi/N,x);
rep(j,N) a[j][i]=x[j];
}
rep(i,N) a[i]=fft(N,-2*pi/N,a[i]);
double ret=0.0,sum=0.0;
rep(i,N) rep(j,N){
if(i>=2*n-1 || j>=2*n-1) continue;
int d=(n-i-1)*(n-i-1)+(n-j-1)*(n-j-1);lint num=(lint)(a[i][j].real()/(N*N)+0.5);
if(num<1L) continue;
me[d]+=num;ret+=sqrt(d+0.0)*num;sum+=num;
}
rep(i,n) rep(j,n){
me[0]-=ma[i][j];sum-=ma[i][j];
}
map<lint,lint>::iterator it=me.begin();
while(it!=me.end()){
if((*it).se>0){
out.pb(mp((*it).fi,(*it).se));
}
it++;
}
sort(All(out));
printf("%.12f\n",ret/sum);
rep(i,min(10000,(int)out.size())) cout<<out[i].fi<<' '<<out[i].se/2<<endl;
}
サイズの上限が1000じゃなくて1024なのがFFTを示唆してるような気がするけど、自分は配列サイズを最初1010とかにしてた所為でWA喰らいました(半ギレ