yukicoder : No.2746 Bicolor Pyramid

これは知らなかった。
https://yukicoder.me/problems/no/2746

問題

青いブロックB個、赤いブロックがR個ある。
ここでk段のピラミッドを作ることを考える。
上からk段目はk^2個の同じ色のブロックを置かなければならない。

最大何段のピラミッドを構築できるか。

解法

f(k) = 1^2+2^2+...+k^2とする。

293段以下の場合、
dp(k,n) := k段のピラミッドを作るときに青いブロックをn個ピッタリ使えるか
というdpテーブルを埋めてみる。
dp(k,n)がTrueかつn≦Bかつn+R≦f(k)ならkが解。

294段以上の場合、B+R≦f(k)となるkが解なので、kを二分探索してよい。

ll A,B;
int ok[295][86400];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>A>>B;
	if(A>B) swap(A,B);
	if(A>=85344) {
		ll v=0;
		while((v+1000)*(v+1001)*(2*v+2001)/6<=A+B) v+=1000;
		while((v+1)*(v+2)*(2*v+3)/6<=A+B) v++;
		cout<<v<<endl;
	}
	else {
		ok[0][0]=1;
		int cand=0;
		ll v=0;
		FOR(i,293) {
			FOR(j,85400) {
				ok[i+1][j]|=ok[i][j];
				if(j>=(i+1)*(i+1)) ok[i+1][j]|=ok[i][j-(i+1)*(i+1)];
			}
			int num=0;
			FOR(j,85400) if(ok[i+1][j]&&j<=A) cand=max(cand,j);
			if(cand+B>=(i+1)*(i+2)*(2*v+3)/6) v=i+1;
		}
		A=min(A,(ll)cand);
		if(v==293) {
			while((v+1000)*(v+1001)*(2*v+2001)/6<=A+B) v+=1000;
			while((v+1)*(v+2)*(2*v+3)/6<=A+B) v++;
		}
		cout<<v<<endl;
	}
}

まとめ

ある程度以上kが大きくなると自由度が上がるの面白いな。