kmjp's blog

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

WUPC 2019 : J - Color Ball

これ系苦手だけど幸いどうにかなった。
https://atcoder.jp/contests/wupc2019/tasks/wupc2019_j

問題

N個の色が異なる筒があり、それぞれA[i]個の筒と同じ色のボールが入っている。
ボールの総数をMとしよう。

ボールをすべて取り出し、並べ替えてまた筒に入れる。
その際、各筒にあるボールの数は最初と同じで、かつ筒と同色のボールは入らないようにしたい。
同じ筒にあるボール群について、色の並びが異なる場合は異なる組み合わせとみなす場合、組み合わせは何通りあるか。

解法

筒が横1列に連結していると考え、ボールをM個並べることを考える。
i番目の筒から順に、その中にあるA[i]個のボールをi番より手前の空きスポットに入れるか、後ろのどこかに入れるかを考えよう。
f(i,n) := i番目までの筒を考えたとき、i番目までのボールのうちn個がすでに入る場所が決まっている場合の並び
を考える。なお、同色のボールの並びも区別するものとする。

この時、f(i,n)から(i+1)番目の筒の状態遷移を考えよう。
S=sum(A[1..n])とする。
(i+1)番目の筒にあるA[i+1]個のボールは、手前に残っている(S-n)個の空きスペースのどこかか、奥に残っているどこかに入る。
また、i番目までの筒のうち入る場所が未確定の(S-n)個のうちいくつかは、(i+1)番目の筒のどこかに入る。
A[i+1]個のボールのうちj個が手前に入り、(S-n)個のうちk個が(i+1)番目の筒に入るとすると、
 f(i+1,n+j+k) += f(i,n)*P(A(i+1),j)*C(S-n,j)*P(A(i+1),k)*C(S-n,k)
となる。

最終的にf(N,0)に対し、同色のボールの並びの区別をしないよう、A[1]!*A[2]!....で割ったものが解。

int N,M;
int A[2020];
int mo=1000000007;

ll dp[2020][2020];
const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
ll comb[2020][2020];
ll perm[2020][2020];

ll L[2020],F[2020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	
	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;
	
	FOR(i,2010) FOR(j,i+1) {
		perm[i][j]=fact[i]*factr[i-j]%mo;
		comb[i][j]=perm[i][j]*factr[j]%mo;
	}
	
	dp[0][0]=1;
	int sum=0,f;
	FOR(i,N) {
		cin>>A[i];
		FOR(x,sum+1) if(dp[x][sum-x]) {
			dp[x][sum-x]%=mo;
			FOR(l,min(A[i],sum-x)+1) L[l]=perm[A[i]][l]*comb[sum-x][l]%mo;
			FOR(l,min(A[i],sum-x)+1) {
				FOR(f,min(A[i],sum-x)+1) {
					if(sum-x+A[i]-(f+l)<=M-(sum+A[i])) (dp[x+l+f][sum-x+A[i]-(f+l)]+=dp[x][sum-x]*L[l]%mo*L[f])%=mo;
				}
			}
		}
		sum+=A[i];
	}
	ll ret=dp[M][0];
	FOR(i,N) ret=ret*factr[A[i]]%mo;
	cout<<ret<<endl;
	
}

まとめ

こっちは自力で解けたけど、時間内には終わらなかった。