AtCoder K2PC Hard A 紅茶

http://k2pc-hard.contest.atcoder.jp/tasks/k2pc001_h1

問題

ab = concatMap (\i-> map (\n-> (i - n, n)) [1..i-1]) [1..]
a = map fst ab
b = map snd ab

という数列がある。
入力i j (1 <= i, j <= 10^8) に対し、(a_i + a_j, b_i, b_j)がabの何番目に出るかを求めろ(ただしインデックスは1-origin, 答えは一つだけ存在する)

回答

まず、任意のi番目のa,bを求めたい。
これには、(a+b)が同じであるサイズがO(N^2)で増えてくことを利用して、O(√N)で適当に求めた。
OEISったら、"floor( 1/2 + sqrt(2*n) )"で出来るらしい。が、long longのfloat関数は怖いし(この大きさなら大丈夫？)、integer版も、2分探索とかすればいいかもしれないが長くなるので、これでいいと思う。
上のhaskellコードを見ればわかりやすいが、逆はそのとおりにやるだけ

コード

#include <string>
#include <vector>
#include <algorithm>
#include <set>
#include <map>
#include <queue>
#include <iostream>
#include <sstream>
#include <cstdio>
#include <cmath>
#include <ctime>
#include <cstring>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define reu(i,l,u) for(int (i)=(int)(l);(i)<(int)(u);++(i))
//#define aut(v, x) __typeof(x) v = (x)
#define aut(v, x) auto v = x
#define each(it,o) for(aut(it, (o).begin()); it != (o).end(); ++ it)
#define all(o) (o).begin(), (o).end()
#define pb(x) push_back(x)
#define mp(x,y) make_pair((x),(y))
#define mset(m,v) memset(m,v,sizeof(m))
using namespace std;
typedef long long ll; typedef vector<int> vi; typedef pair<int,int> pii; typedef vector<pii> vpii;
const int INF = 0x3f3f3f3f; const ll INFL = 0x3f3f3f3f3f3f3f3fLL;
 
int main() {
	int ii,jj;
	scanf("%d%d",&ii,&jj);
	int i = 1, x = 1;
	int a, b;
	while(1){
		if(ii < x + i) {
			b = ii - x + 1; a = i;
			break;
		}
		x += i;
		i ++;
	}
	i = 1, x = 1;
	while(1){
		if(jj < x + i) {
			b += jj - x + 1; a += i;
			break;
		}
		x += i;
		i ++;
	}
//	cerr << a << ", " << b << endl;
	ll r = 0;
	rer(i, 1, a) r += i;
	r += b;
	cout << r << endl;
	return 0;
}