kmjp's blog

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

AtCoder ABC #156 : E - Roaming

今回は普通の?傾斜だったっぽい。
https://atcoder.jp/contests/abc156/tasks/abc156_e

問題

N個の部屋があり、最初それぞれ1人ずつ人がいる。
どこかの部屋にいる人を1人選び、異なる部屋に動かす、ということをK回行う。
人は互いに区別しないとき、移動後の人数分布は何通りか。

解法

N≦Kなら、最後N回で各人所望の場所に動かすことで任意の配置にできるのでH(N,N)。
K<Nの場合を考える。
0人となる場所がm箇所あるとする。そのm人は残り(N-m)部屋に散らばるのでComb(N,m)*H(N-m,m)通り。
mを0~Kまで総当たりしよう。ただしK=1の時m=1の一択なので注意。

int N,K;
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(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>>K;
	
	if(K>=N) {
		cout<<hcomb(N,N)<<endl;
		return;
	}
	if(K==1) {
		cout<<1LL*N*(N-1)%mo<<endl;
	}
	else {
		ll ret=0;
		FOR(i,K+1) {
			int lef=N-i;
			ret+=comb(N,i) * hcomb(lef,i)%mo;
		}
		cout<<ret%mo<<endl;
	}
	
}

まとめ

ABCの問題名、Writerによるんだろうけどよくわからん命名規則だな…。