読者です 読者をやめる 読者になる 読者になる

本履歴

購入した古本の履歴と時々プログラミング

CodeChefネタ

今月(http://www.codechef.com/NOV11)は、New Restaurant(http://www.codechef.com/NOV11/problems/NEWREST)が考えさせられた。

問題文はこんな感じ

自然数N,M,Kが与えられたとする。M個の異なった要素のものをN個重複順列する中で要素の種類が高々K個になる並べ方の総数は?
でかい数になるので、1000000007で割った余りを求めてね。これをT回問いますよ。

制限
1 ≤ T ≤ 100
1 ≤ N ≤ 1000
1 ≤ M ≤ 1000000
1 ≤ K ≤ 1000

最初考えたのは、M個の中からK個まず選んでおいて、それをN個並べようというもの。これなら高々Kを満たすことができる。

\left(\large\begin{array}{GC+23}M\\K\end{array}\right)k^n

でも、これだと同じものを重複して数えてしまうね。
例えば、M,N,K=3,2,2の場合、{A,B,C}から2個{A,B}と選んで2個重複順列すると{AA,AB,BA,BB}、
同じように{A,C}と選ぶと{AA,AC,CA,CC}で、AAがダブってる。
色々考えてみて、高々K個というのが曲者と気づいた。そこで、ちょうどK個の総数を考えてみた。Inclusion-Exclusion原理で、

S(N,M,K)=\left(\large\begin{array}{GC+23}M\\K\end{array}\right)\left{ K^N-\left(\large\begin{array}{GC+23}K\\1\end{array}\right)(K-1)^N+...+(-1)^K\left(\large\begin{array}{GC+23}K\\K\end{array}\right)(K-K)^N\right}

従って求めるものは、

\sum_{l=1}^{K}S(N,M,l)

さて、この式をにらむと、同じ項がたくさん出てきているので、事前に計算しておけば、O(T*K*K)で計算できるはず・・・

l^N, \hspace{5}\left(\large\begin{array}{GC+23}l\\t\end{array}\right), \hspace{5}l^{-1}, \hspace{15} 1\leq t \leq l \leq K

次のようなコードが出来た。

{$R-,S-,Q-,I-,O+}
const MM=1000000007;
	MMMM=MM*MM;

var t,n,k,l,j:longword;m,ans,temp,temp1,temp2,coe:qword;
	sign:boolean;
	comibination:array[0..1005,0..1005] of qword; //the binomial coefficient indexed by n,k
	power:array[1..1005] of qword; //power[l]=l^n%MM
	inverse:array[1..1005] of qword; // inverse[n]=n^(-1) over Z_MM

procedure init_comb;
var i,j:longword;
begin
	for i:=0 to 1000 do comibination[i,0]:=1;
	for i:=1 to 1000 do for j:=1 to i do comibination[i,j]:=(comibination[i-1,j]+comibination[i-1,j-1]) mod MM;
end;

function inv(a:longword):qword;
var d0,d1,d2,x0,x1,x2,q:int64;
begin
	d0:=int64(a);d1:=int64(MM);x0:=1;x1:=0;
	while d1<>0 do begin
		q:=d0 div d1;
		d2:=d0-q*d1;
		x2:=x0-q*x1;
		x0:=x1;
		x1:=x2;
		d0:=d1;
		d1:=d2;
	end;
	inv:=qword((x0+MMMM) mod MM);
end;

procedure init_inverse;
var i:longword;
begin
	for i:=1 to 1000 do inverse[i]:=inv(i);
end;

//a^n%MM
function powmod(a:qword;n:longword):qword;
var prod:qword;
begin
	prod:=1;
	while n>0 do begin
		if odd(n) then prod:=(prod*a) mod MM;
		n:=n shr 1;
		a:=(a*a) mod MM;
	end;
	powmod:=prod;
end;

procedure init_pow(k,n:longword);
var i:longword;
begin
	for i:=1 to k do power[i]:=powmod(qword(i),n);
end;

begin
	init_comb;
	init_inverse;
	readln(t);
	while t>0 do begin
		readln(n,m,k);
		init_pow(k,n);//writeln('inv done!');
		ans:=m;
		coe:=m;
		for l:=2 to k do begin
			temp1:=power[l];
			temp2:=0;
			sign:=true;
			for j:=1 to l-1 do begin
				sign:=not sign;
				if sign then
					temp1:=(comibination[l,j]*power[l-j]+temp1) mod MM
				else //negative terms
					temp2:=(comibination[l,j]*power[l-j]+temp2) mod MM;
			end;
			temp:=(temp1+MMMM-temp2) mod MM;
			coe:=(coe*(m-l+1) mod MM)*inverse[l] mod MM;
			ans:=(coe*temp+ans) mod MM;
		end;
		writeln(ans);
		dec(t);
	end;
end.

ideoneで実行させていたのだが、どうしても2秒近くかかってしまう。ideoneはcodechefよりだいたい3倍速いので、余裕でTLEですね。
式を睨んでみたが、交代和が入った式が簡単にはならなさそう。うーん行き詰まった。上位の人は0.2secで解いているし、他にエレガントな解があるに違いない。
そこで、自分が何を求めているのかもう一度考えてみた。ちょうどK個ということは、NをK個のグループに分けていることだな。
これ、どっかで聞いたなー。あっ、wikipedia:スターリング数だ!
ということで、まずM個からK個選んでおいて、NをK個のグループに分割してそのグループにK個を当てはめていけばOK。

S(N,M,K)=\left(\large\begin{array}{GC+23}M\\K\end{array}\right)K!\left{\large\begin{array}{GC+23}N\\K\end{array}\right}

簡単にして、P(M,K)を相違なるM個からK個選んで並べる順列の総数とすると、

S(N,M,K)=P(M,K)\left{\large\begin{array}{GC+23}N\\K\end{array}\right}

スターリング数は事前に計算させておいて、コードに落とせば、

{$R-,S-,Q-,I-,O+}
const MM=1000000007;

var t,n,m,k:longword;
	s:array[0..1005,0..1005] of qword; //Stirling number of the second kind

procedure init_stirling;
var i,j:longword;
begin
	for i:=1 to 1000 do s[i,1]:=1;
	for i:=2 to 1000 do for j:=2 to i do s[i,j]:=(s[i-1,j-1]+j*s[i-1,j]) mod MM;
end;

function ans(n,m,k:longword):qword;
var r,p:qword;i:longword;
begin
	r:=0;
	p:=1;
	for i:=1 to k do begin
		p:=(p*(m-i+1)) mod MM;
		r:=(r+p*s[n,i]) mod MM;
	end;
	ans:=r;
end;

begin
	init_stirling;
	readln(t);
	while t>0 do begin
		readln(n,m,k);
		writeln(ans(n,m,k));
		dec(t);
	end;
end.

O(T*K+K*K)のちょうシンプルな34行プログラムになって、無事Accepted!

結局、あの交代和はスターリング数だったのだねぇ。へえ
あと、ちょっとでも速くしようとレンジチェックを外すコマンドラインスイッチを入れているが、
これはlocal or ideoneテストの時は外したほうがいいね。桁あふれのバグを見逃す危険性が大って、当たり前だなあ。

今年初めにたてた目標「毎月CodeChefを一問は解く」、3月以外はなんとか皆勤。来月もがんばろう。