kmjp's blog

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

第6回 ドワンゴからの挑戦状 予選 : C - Cookie Distribution

これは問題文の言い換え方を覚えていたので何とか解けたけど、苦労しすぎた。
https://atcoder.jp/contests/dwacon6th-prelims/tasks/dwacon6th_prelims_c

問題

N人の子供がいる。
彼らにK回クッキーを配る。
i回目ではN人中A[i]人に1個ずつクッキーを与える。

クッキーの配り方全通りにおいて、各人の得たクッキー数の積の総和を求めよ。

解法

問題文を言い換える。
N行K列のグリッドがあったとする。
クッキーを配ることを、i列目のうちA[i]行個のマスを黒く塗ったと言い換える。

この状態で、クッキー数の積の総和は以下のように言い換えられる。

  • 各人がいずれかの回次でクッキーを割り当てられているかどうかを判定したとき、全員が対象の回次でクッキーが割り当てられているような配り方の判定法および配り方を数え上げよ。

こうすると、
f(i,n) := i回目までの回次で、割り当て未確定の人がn人となる配り方の総和
とすると、f(0,N)=1から初めてf(K,0)を求めるとそれが解になる。
f(i,n)の次の遷移は、残されたn人中m人が(i+1)回目の回次で数え上げ対象になったとすると、
f(i+1,n-m) += f(i,n) * Comb(n,m) * Comb(N-m, A[i]-m)
で数え上げられる。
1つ目のCombinationは(i+1)回目の回次を数え上げ対象にするためn人中m人を選ぶもので、2つ目のCombinationはこのm人以外のクッキーの割り当て型である。

int N,K;
int A[22];
const ll mo=1000000007;
const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll comb(ll N_, ll C_) {
	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 from[1010];
ll to[1010];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	comb(1,1);
	
	int sum=0;
	FOR(i,K) {
		cin>>A[i];
		sum+=A[i];
	}
	
	ll ret=0;
	from[N]=1;
	FOR(i,K) {
		ZERO(to);
		FOR(x,N+1) if(from[x]) {
			FOR(y,A[i]+1) if(x>=y) {
				(to[x-y]+=from[x]*comb(x,y)%mo*comb(N-y,A[i]-y))%=mo;
			}
		}
		swap(from,to);
	}
	
	cout<<from[0]<<endl;
	
	
	
}

まとめ

積の形をこのグリッドの数え上げに言い換えるテクを覚えていたのでよかった。