kmjp's blog

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

yukicoder : No.802 だいたい等差数列

未だに包除原理が漢字変換できません。
https://yukicoder.me/problems/no/802

問題

整数N,M,D1,D2が与えられる。
以下を満たす整数列Aは何通りか。

  • AはN要素の単調増加列で、1~Mの範囲を取る。
  • 隣接要素間の差はD1以上D2以下である。

解法

B[i] = A[i]-D1*(i-1)-1 とする。
するとBの隣接要素間の差は0以上(D2-D1)以下であり、かつ0~(M-(D2-D1)*(N-1))-1の範囲の値を取ることになる。

D=D2-D1、L=(M-(D2-D1)*(N-1))-1とおくと、この問題は0~Lの範囲を取るN要素の単調増加列Bは何通りか求める問題となる。
仮にB[0]=0、B[N+1]=Lと固定し、B[1]~B[N]を求めることを考えよう。
B[0]~B[N+1]の間には差分が(N+1)個あるので、差がD以下の条件を外せば、解は \displaystyle {}_{N+1} H_Lである。

B[1]~B[N]の間にある(N-1)個の差分のうち、(D+1)以上差が生じるケースを考えて包除原理を適用することを考える。
差分のうちi個が(D+1)以上の差が生じるとすると
 \displaystyle \sum_{i=0}^{N-1} ( (-1)^i \times {}_{N+1} H_{L-i(D+1)} \times {}_{N-1} C_i )
が解となる。

ll N,M,D1,D2,D;
ll mo=1000000007;

ll comb(ll N_, ll C_) {
	const int NUM_=1400001;
	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_);}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>D1>>D2;
	D=D2-D1;
	M-=1LL*(N-1)*D1;
	M--;
	
	ll ret=0;
	for(i=0;i<=N-1;i++) {
		if(1LL*i*(D+1)>M) continue;
		ll a=hcomb(N+1,M-(i*(D+1)))*comb(N-1,i)%mo;
		if(i%2==0) ret+=a;
		else ret+=mo-a;
	}
	cout<<ret%mo<<endl;
	
	
}

まとめ

最初O(NM)解ができて、これ包除原理で計算量落とせるな…と思ったところで離脱。