Codeforces Round #273 (Div. 2) D. Red-Green Towers

問題

r個の赤いブロックとg個の緑色のブロックで、

1段目にn個、
2段目にn-1個、
3段目にn-2個、
・・・
N段目に1個、という三角形型のオブジェを作り上げます。
ただし、それぞれの段は全て同じ色のブロックで作らなければなりません。

　　　赤
　　緑・緑
　緑・緑・緑
赤・赤・赤・赤

例えばr=g=5のときはこんな感じのオブジェが作れます。
両方のブロック合わせて作れるオブジェの最高段数をhとしたとき、
h段のオブジェの作り方の場合の数（mod 1000000007）を求めなさい。

制約

0<=r,g<=200000
1<=r+g

解法

N段作るのに必要なブロック数はN(N+1)/2なので、
r,gが最大の値を取っても、最大の段数は800段程度です。
なので、O(h*(r+g))くらいがギリギリ間に合いそうな感じです。

また、それぞれの段の作り方は2通りしかないため、
片方の色のブロックについてだけどの段に使うかを考えれば、もう片方をどこに使うのかも同時に考えられたことになります。

というわけで、赤ブロックをx個使った場合、緑ブロックはN(N+1)/2-x個使うことになります。
この「赤いブロックをx個使う」という場合の数の総和が求めるべき答えとなります。

例えば、上の具体例では赤ブロックを5個使っていますが、そのような使い方は
・1段目と4段目
・2段目と3段目
の2通りがあります。なのでr=g=5のときは作り方は全部で2通りです。

もう少し例を増やすと、赤ブロックを6個使う使い方は
・1段目と5段目
・2段目と4段目
・1段目と2段目と3段目
で3通りです。

この場合の数は

dp[0]=1
for(i in 1..n)for(j in m..0)dp[i+j]+=dp[j]

という簡単なdpで求まります（iが「一番最後に足す数」でjが「足す前の総和」という感じです）。

ところで、r=g=5のとき、赤ブロックを4個しか使わなかった場合、
h=4のオブジェを作るには緑ブロックを6個使わないといけませんが、足りません。
ということで、最後足し算するときにどちらかのブロックの個数が足りてないような組み合わせは足さないようにする必要があります。

プログラム

#include<bits/stdc++.h>
using namespace std;


#define li			long long int
#define rep(i,to)	for(li i=0;i<((li)(to));i++)
#define repp(i,start,to)	for(li i=(li)(start);i<((li)(to));i++)
#define pb			push_back
#define sz(v)		((li)(v).size())
#define bgn(v)		((v).begin())
#define eend(v)		((v).end())
#define allof(v)	(v).begin(), (v).end()
#define dodp(v,n)		memset(v,(li)n,sizeof(v))
#define bit(n)		(1ll<<(li)(n))
#define mp(a,b)		make_pair(a,b)
#define rin	rep(i,n)
#define EPS 1e-10
#define ETOL 1e-8
#define MOD 1000000007
typedef pair<li, li> PI;

#define INF bit(60)

#define DBGP 1


#define idp if(DBGP)
#define F first
#define S second
#define p2(a,b)		idp cout<<a<<"\t"<<b<<endl
#define p3(a,b,c)		idp cout<<a<<"\t"<<b<<"\t"<<c<<endl
#define p4(a,b,c,d)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<endl
#define p5(a,b,c,d,e)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<endl
#define p6(a,b,c,d,e,f)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<"\t"<<f<<endl
#define p7(a,b,c,d,e,f,g)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<"\t"<<f<<"\t"<<g<<endl
#define p8(a,b,c,d,e,f,g,h)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<"\t"<<f<<"\t"<<g<<"\t"<<h<<endl
#define p9(a,b,c,d,e,f,g,h,i)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<"\t"<<f<<"\t"<<g<<"\t"<<h<<"\t"<<i<<endl
#define p10(a,b,c,d,e,f,g,h,i,j)		idp cout<<a<<"\t"<<b<<"\t"<<c<<"\t"<<d<<"\t"<<e<<"\t"<<f<<"\t"<<g<<"\t"<<h<<"\t"<<i<<"\t"<<j<<endl
#define foreach(it,v)	for(__typeof((v).begin()) it=(v).begin(); it!=(v).end(); ++it)
#define p2p(x)		idp p2((x).F, (x).S)
#define dump(x,n)	idp{rep(i,n){cout<<x[i]<<" ";}puts("");}
#define dump2(x,n)	idp{rep(i,n){cout<<"["<<x[i].F<<" , "<<x[i].S<<"] ";}puts("");}
#define dumpi(x)	idp{foreach(it, x){cout<<(*it)<<" ";}puts("");}
#define dumpi2(x)	idp{foreach(it, x){cout<<"["<<(it)->F<<" , "<<(it)->S<<"] ";}puts("");}

#define read2d(a,w,h)	rep(i,h)rep(j,w)cin>>a[i][j]
#define dump2d(a,w,h)	rep(i,h){rep(j,w)cout<<a[i][j]<<" ";puts("");}

typedef pair<li, li> PI;

inline li num(li x){
	return (x*x+x)/2;
}

li caln(li x){
	li l=0, r=1000000000;
	while(l+1<r){
		li mid=(l+r)/2;
		if(num(mid)>x)r=mid;
		else l=mid;
	}
	return l;
}

li dp[200200];

int main(){
	li cube[2];
	rep(i,2)cin>>cube[i];
	sort(cube, cube+2);

	li n=caln(cube[0]+cube[1]);

	rep(i,cube[0]+1)dp[i]=0;
	dp[0]=1;

	repp(i,1,n+1){
		for(li j=cube[0]-i; j>=0; j--){
			dp[j+i]=(dp[j]+dp[j+i])%MOD;
		}
	}

	li res=0;
	li nn=num(n);
	rep(i,cube[0]+1){
		if(dp[i]>0 && cube[0]>=i && (nn-i)<=cube[1])res=(res+dp[i])%MOD;
	}
	cout<<res<<endl;

	return 0;
}

ちなみに、hはr+gのみに依存します。
つまりr=1, g=9だろうがr=g=5だろうが、r+g=10なのでh=4は確定、ということです。
0<=x<=N(N+1)/2　なるxを{1..N}の0個以上の部分和から作る方法は必ず1つ以上存在する、というのがその根拠ですが、
感覚的な証明もどきとしては、いくつかの数（集合Sとします）の和でxが表現できているとき、
x+1やx-1も簡単に作ることが出来る、ということです。
Sに1が含まれない→S'=S⋃{1}
Sに1が含まれる→Sの中のどれかの要素を+1する
という感じでsum(S)=xなるSからsum(S')=x+1なるS'を作ることが出来ます。

具体的にはS={2,4}のとき
{2,4}->{1,2,4}->{1,3,4}
という感じに和を6,7,8と増やしていくことが可能です。