skyaozoraの日記

AOJ2560 Point Distance

23:10

問題概要

問題文

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喰らいました（半ギレ