AtCoder AGC #002 : F - Leftmost Ball

これは本番自力では思いつかなさそう。
http://agc002.contest.atcoder.jp/tasks/agc002_f

問題

1～NのN色のボールがそれぞれK個ずつある。
これらを横１列に並べたうえ、各色一番左のボールを色0で塗り替える。
こうして得られるボール列は何通りあるか。

解法

塗り替えられてないボールは左から色1～Nの順で登場すると仮定しよう。
(実際はそうではないが、問題の対称性よりその分最後にN!を掛算すればつじつまが合う)

以下のような有向グラフを考える。

• 頂点が横N列×縦K行で並んでいる。
• 各頂点すぐ下の点に有向辺がある。
• １行目と２行目について、各頂点すぐ右の点に有向辺がある。

このグラフのトポロジカルソート順の組み合わせを求めれば(そしてそこに最初に書いたN!を掛ければ)それが解。
r行c列の頂点を辿ることは、色rのボールのc個目を置くことに相当する。

以下のDPを考える。
dp(x,y) := 最上位行がx個、2行目以降がy個未到達の頂点が残った状態からのトポロジカルソートの組み合わせ
dp(N,N)が求める解となる。

• x==yの時、最初到達可能な頂点は最上位左端の頂点一択なので、dp(x,y)=dp(x-1,y)
• x＜yの場合、最上位行左端の頂点を辿る場合と、２行目左端の頂点を辿るケースがある。
  • 前者の場合、残った状態はdp(x-1,y)
  • 後者の場合、残った状態はdp(x,y-1)である。また今回頂点を辿った列には(K-2)個頂点が残る
  • よって両者を合わせて

ll mo=1000000007;
ll combi(ll N_, ll C_) {
	const int NUM_=4500001;
	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;
}

int N,K;
ll dp[2020][2020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	if(K==1) return _P("1\n");
	
	dp[0][1]=1;
	for(y=2;y<=N;y++) dp[0][y] = dp[0][y-1]*combi(y*(K-1)-1,K-2)%mo;
	
	for(x=1;x<=N;x++) {
		
		if(x==1) dp[1][1]=1;
		else dp[x][x]=dp[x-1][x];
		
		for(y=x+1;y<=N;y++) dp[x][y] = (dp[x-1][y] + dp[x][y-1]*combi(x+y*(K-1)-1,K-2))%mo;
	}
	
	ll ret=dp[N][N];
	for(i=1;i<=N;i++) ret=ret*i%mo;
	cout<<ret<<endl;
}

まとめ

なぜ色を上塗りしてしまうのか。