yukicoder : No.2368 I love a square root of 2

Last Problemを思い出すな…。
https://yukicoder.me/problems/no/2368

問題

非負整数a,bを用いて、a+√2bの形で表せる実数のうち、小さい順にN番目のものは何か。

解法

まず解の整数部分Xを二分探索しよう。
aをO(√N)程度まで総当たりすれば、それぞれ取りえるbの個数はすぐ求められるので、O(√N*logN)で二分探索できる。

整数部分が定まったら、X以上(X+1)未満のa+√2bを列挙しよう。
これをソートするわけだが、a+√2bとc+√2dの大小は、(a-c)と(b-d)√2の符号を比較し、符号が一緒なら2乗を比較することで整数の範囲で誤差なく判定できる。

ll N;

ll hoge(int V) { //V以下の値
	ll ret=0;
	
	for(int a=0;a<=min(V,1000000);a++) {
		ll d=V-a;
		ret++;
		ret+=d/sqrt(2);
	}
	return ret;
}

bool cmp(pair<int,int> L,pair<int,int> R) {
	ll a=L.first-R.first;
	ll b=R.second-L.second;
	assert(a!=0&&b!=0);
	if(a<0&&b>0) return 1;
	if(a>0&&b<0) return 0;
	if(a<0&&b<0) return a*a>2*b*b;
	return a*a<2*b*b;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	if(N==1) {
		cout<<"0 0"<<endl;
		return;
	}
	int L=0,R=1000000;
	while(R-L>1) {
		int M=(L+R)/2;
		if(hoge(M)>=N) {
			R=M;
		}
		else {
			L=M;
		}
	}
	
	vector<pair<int,int>> V;
	N-=hoge(L);
	for(i=0;i<R;i++) {
		int na=(R-i)/sqrt(2);
		int nb=(L-i)/sqrt(2);
		if(na>nb) V.push_back({i,na});
	}
	V.push_back({R,0});
	assert(V.size()==hoge(R)-hoge(L));
	sort(ALL(V),cmp);
	cout<<V[N-1].first<<" "<<V[N-1].second<<endl;
}

まとめ

最初符号を考慮し忘れて焦った。