TopCoderOpen 2019 Round3A : Easy FamilySeatingArrangement

Easyだけでレートが上がるとは。
https://community.topcoder.com/stat?c=problem_statement&pm=15570&rd=17611

問題

N家族がK個のテーブルに分かれて座る。
各家族iは、親が1名と、子供がA[i]名からなる。
子供は、自分の親と同じテーブルか、誰の親もいないテーブルにのみ座ることができる。

人やテーブルを区別できるとき、テーブルに座る人の組み合わせは何通りか。

解法

自分の問題に近い。テーブルを区別する点が異なるけど。
yukicoder : No.140 みんなで旅行 - kmjp's blog

まず、親の配置を決めよう。
これはテーブルを区別する点を除けばスターリング数と同じ感じでO(N^2)のDPで解ける。

親がm個のテーブルに分かれた場合を考える。
各子供の座る選択肢は親と同じか誰の親もいないところ、すなわち(k-m+1)通りの選択肢がある。

ll dp[1010][1010];
ll mo=1000000007;

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;
}

class FamilySeatingArrangement {
	public:
	int countWays(vector <int> a, int k) {
		int N=a.size();
		ZERO(dp);
		int i,j;
		dp[0][0]=1;
		FOR(i,N) {
			FOR(j,N) {
				(dp[i+1][j]+=dp[i][j]*j)%=mo;
				if(j<k) (dp[i+1][j+1]+=dp[i][j]*(k-j))%=mo;
			}
		}
		
		ll ret=0;
		for(i=1;i<=k;i++) {
			int ok=1+(k-i);
			ll pat=dp[N][i];
			FOR(j,N) (pat*=modpow(ok,a[j]))%=mo;
			ret+=pat;
		}
		return ret%mo;
		
	}
}

まとめ

最初テーブル区別しないと勘違いしてタイムロス。