kmjp's blog

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

第一回日本最強プログラマー学生選手権-予選- : F - Candy Retribution

あと一歩だった気がするけどちょっと及ばず。
https://atcoder.jp/contests/jsc2019-qual/tasks/jsc2019_qual_f

問題

整数N,M,L,Rが与えられる。
以下を満たすN要素の整数列Aは何通りか。

  • Aの各要素は非負整数
  • Aの要素を降順に並べたとき、M番目と(M+1)番目の値は等しい
  • Aの総和はL以上R以下

解法

2つ目の条件を無視すると、組み合わせはH(N,L)+H(N,L+1)+...+H(N,R)である。
ここから2つ目の条件に反するものを引こう。
この場合、M番目と(M+1)番目が異なるケースを引くことを考える。

M番目の値をちょうどaとし、(M+1)番目以降がa未満となるケースを考えよう。
F(a,b,x) := M番目の値がa以上で、(M+1)番目以降の値がb未満で、総和がx以下となる組み合わせ
とする。M番目の値が(a+1)以上のケースを無視し、また総和はL以上R以下でないといけないので、aを総当たりしながら
(F(a,a,R)-F(a+1,a,R))-(F(a,a,L-1)-F(a+1,a,L-1))
の総和を求めればよい。

あとはF(a,b,x)の求めかたである。
まずN個中M個はa以上でなければならないので、C(N,M)個だけa以上となる要素を決め、先に固定でa*M個分の値を確保して割り振ったと考えよう。
あとは、残り(N-M)個がb未満でなければならないので、ここでb個以上になる個数を列挙しつつ、包除原理の要領で、求めよう。
仮にb個以上になる要素がp個とする。するとAの要素のうちa*M+b*p個分は割り当てが決まるので、残りを(N+1)人で割り振ることを考えると
C(N-M,p)*H(N+1,x-(a*M+b*p))をpの偶奇に応じて足し引きしていくとよい。

ll N,M,L,R;
ll mo=1000000007;
ll ret;

ll comb(ll N_, ll C_) {
	const int NUM_=700001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}


ll pat(int a,int b,int ma) {
	ll tot=0;
	for(int over=0;over*b+M*a<=ma && over<=N-M;over++) {
		ll pat=comb(N-M,over)%mo*hcomb(N+1,ma-(over*b+M*a))%mo;
		if(over%2==0) {
			tot+=pat;
		}
		else {
			tot+=mo-pat;
		}
	}
	return (tot%mo+mo)%mo*comb(N,M)%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>L>>R;
	
	for(i=L;i<=R;i++) ret+=hcomb(N,i);
	for(int la=1;la<=R;la++) {
		ret-=(pat(la,la,R)-pat(la+1,la,R))-(pat(la,la,L-1)-pat(la+1,la,L-1));
	}
	
	
	cout<<(ret%mo+mo)%mo<<endl;
	
}

まとめ

うーむ、これ系苦手だな…。