kmjp's blog

競技プログラミング参加記です

AtCoder ABC #234 : Ex - Enumerate Pairs

計算量の見積もりが本質?
https://atcoder.jp/contests/abc234/tasks/abc234_h

問題

2次元座標上でN個の点の座標が指定される。
頂点対のうち、距離がK以下のものを列挙せよ。
なお、解はM個以下である入力しか与えられない。

解法

領域をK*Kの矩形のグリッドに区切り、各頂点について、自身の属するセルと8方向の隣接セル内の頂点だけ総当たりすればよい。
「解がM個以下」の条件がないと、最悪O(N^2)かかりTLEするが、「解がM個以下」の条件により上記総当たりがO(N+M)で抑えられる。証明はEditorial参照。

int N;
ll K;
ll X[202020],Y[202020];

map<pair<int,int>,vector<int>> cand;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	vector<pair<int,int>> ret;
	FOR(i,N) {
		cin>>X[i]>>Y[i];
		x=X[i]/K;
		y=Y[i]/K;
		for(int a=x-1;a<=x+1;a++) {
			for(int b=y-1;b<=y+1;b++) {
				if(cand.count({a,b})) {
					vector<int> c=cand[{a,b}];
					FORR(j,c) {
						if((X[i]-X[j])*(X[i]-X[j])+(Y[i]-Y[j])*(Y[i]-Y[j])<=K*K) ret.push_back({j+1,i+1});
					}
				}
			}
		}
		cand[{x,y}].push_back(i);
	}
	sort(ALL(ret));
	cout<<ret.size()<<endl;
	FORR2(a,b,ret) cout<<a<<" "<<b<<endl;
}

まとめ

本番バケットに区切ることは考えたものの、O(N^2)かかると思って試さなかった…。