kmjp's blog

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

AtCoder ARC #110 (鹿島建設プログラミングコンテスト2020) : D - Binomial Coefficient is Fun

これがすんなり解けてよかった。
https://atcoder.jp/contests/arc110/tasks/arc110_d

問題

N要素の整数列Aが与えられる。
また、N要素で和がM以下の整数列Bを考える。
全Bの取り方における、Prod(Comb(B[i],A[i]))を求めよ。

解法

「M以下」の条件が面倒なので、末尾にA[N+1]=0を追加して考えよう。
B[1]~B[N]の和がM未満でも、残りをB[N+1]に押し付ければよいことになる。

以降、M以下でなくSum(B)=Mのケースを考える。
M個中、各Comb(B[i],A[i])で選出されるSum(A)個の要素を固定しよう。
残り(M-Sum(A))個の要素をどこに配置するかを考える。

B[i]-A[i]個の要素は、A[i]+1個の隙間に配分していくことになる。
そこで、解はH(sum(A[i]+1),M-sum(A))となる。

int N,M;
int A[2020];

const ll mo=1000000007;

ll modpow(ll a, ll n=mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll comb(int P_,int Q_) {
	if(P_<0 || Q_<0 || Q_>P_) return 0;
	ll p=1,q=1;
	Q_=min(Q_,P_-Q_);
	for(int i=1;i<=Q_;i++) p=p*P_%mo, q=q*i%mo,P_--;
	
	return p*modpow(q,mo-2)%mo;
}
ll hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);};

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	ll space=0,sum=0;
	FOR(i,N) {
		cin>>A[i];
		space+=A[i]+1;
		sum+=A[i];
	}
	space++;
	cout<<hcomb(space,M-sum)<<endl;
}

まとめ

組み合わせ系苦手なので、こういうのがいつもすんなりできるといいんだけどね。