kmjp's blog

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

ABBYY Cup 3.0 : D. PE Lesson

シンプルながらも良い問題。
http://codeforces.com/contest/316/problem/D3

問題

N人の生徒が1列に並んでおり、それぞれ順に1~N番のボールを持っている。
各生徒は、他の生徒とボールを交換できる。

各生徒の最大交換可能回数が与えられたとき、最終的に得られるボール配置の組み合わせ数を答えよ。

解法

各生徒は1回か2回しかボールを投げられない。
まずは1回しかボールを投げられない生徒がa人いるとし、交換の組み合わせ数をI(a)とする。

ある生徒のアプローチは以下のどちらか:

  • 誰とも交換しない場合、他の(a-1)人の組み合わせI(a-1)通り
  • 他の(a-1)人と交換する場合、残された(a-2)人の組み合わせから(a-1)*I(a-2)通り

よってI(a) = I(a-1) + (a-1)*I(a-2)となる。
先にa人の交換相手を決めた後で、2回ボールを投げられる生徒b人の交換先を決める。
b人中i番目の生徒は、すでに決まっているa+(i-1)人の生徒のうちから交換先を決めるか、交換しないことができるので(a+i)通りの選択肢ができる。

よってI(a)に\prod^b_{i=1} (a+i) = \frac{(a+b)!}{a!}を掛ければよい。
O(a+b)で完了する。

int N,P[3];
ll mo=1000000007;
ll I[1000001];

void solve() {
	int f,r,i,j,k,l, x,y;
	
	cin>>N;
	FOR(i,N) P[GETi()]++;
	
	I[0]=I[1]=1;
	for(i=2;i<=1000000;i++) I[i]=(I[i-1]+(i-1)*I[i-2]) % mo;
	ll ret = I[P[1]];
	for(i=P[1]+1;i<=P[1]+P[2];i++) ret = (ret*i) % mo;
	cout << ret << endl;
	
	return;
}

まとめ

コードはかなりシンプルに書けていいね。