Mujin Programming Challenge 2018 F - チーム分け
考え方を知らないと結構難しい。
知ってしまえばこの手の解法は結構すぐに見えるようになると思う。
問題
$n$人について、$i$人目の人は$a_i$人以下からなるチームに所属したいという制約が与えられた元で、チーム分けの場合の数を求めよ。
考え方
まず、DP(or 包除)ぐらいしか解く方法がなさそう。
非自明なDPの考察をするときは、どの状態をまとめるかがポイント。
$a_i$が小さい方が制約が強いので、$a_i$が大きい方から見ていくと、新たに追加する人の制約のみ考えればよい(これまでの人の制約は気にしなくてよい)こともなんとなくわかる。
ここまで読んでも分からなかったら、この問題は一度stackに積んで以下の問題を考えてみよう。
ARC067 E - Grouping
問題
$N$人の人を、$A$人以上$B$人以下のグループに、$i$人のグループが$0$個または$C$個以上$D$個以下になるように分ける方法は何通りあるか。
考え方
DPで考える。 $i$人のグループの個数に関する制約があるので、それを状態としてまとめればよさそう。
$A$人のグループ、$A+1$人のグループ、…、$B$人のグループと順番に(同じ人数のグループは同時に決める)決めていき、 $$ \mathrm{dp}[x][y] = x人のグループを決める段階でy人余っている場合の数 $$ を更新していく。 更新式は、$k = \{0, C, C + 1, C + 2, …, D\}$に対して $$ \mathrm{dp}[x + 1][y - k * x] \mathrel{+}= \mathrm{dp}[x][y] * \mathrm{choose} * \mathrm{arrange} $$ where $$ \mathrm{choose}(y人からkx人選ぶ): {}_y C_{kx} \\\ \mathrm{arrange}(kx人をx人kチームに分ける): \frac{(kx)!}{(x!)^k * k!} $$ とすればよい($k = 0$も条件を満たすことに注意)。
一見$\mathcal{O}(n^3)$に見えるが、$k$の範囲を考えると$x, k$のループで調和級数になっていて、全体で$\mathcal{O}(n^2 \log(n))$。
実装
powの部分をサボっているので$\mathcal{O}(\log(n))$ぐらい余計についています。
// mint, fact, fact_invは省略
signed main() {
int n, A, B, C, D; cin >> n >> A >> B >> C >> D;
init_factorial(n);
vector<vector<mint>> dp(n + 2, vector<mint>(n + 1));
dp[A][n] = 1;
REPI (x, A, B + 1) {
REP (y, n + 1) {
for (int k = 0; k <= D; k = k == 0 ? C : k + 1) {
if (y - k * x < 0) break;
mint choose = comb(y, k * x);
mint arrange = fact[k * x] * pow(fact_inv[x], k) * fact_inv[k];
dp[x + 1][y - k * x] += dp[x][y] * choose * arrange;
}
}
}
cout << dp[B + 1][0] << endl;
}
stack.pop()
さて、前の問題の考察から、グループ分けの場合の数を数える問題では「グループの人数ごとにまとめて考える」ことがポイントであることがわかった。
今回の問題も全く同じ考え方を用いればよい。
一人一人に制約がついているのが厄介だが、少し考えると、大きい順に見ていってそのときに見ているものの制約だけ気にすればいいことがわかる。
というわけで、 $$ \mathrm{num}[x]: a_i = xであるような人の数 $$ を計算しておく。大きい順にグループを決めていき、 $$ \mathrm{dp}[x][y]: x人のグループを作るときにy人余っている場合の数 $$ というDPを、 $$ \mathrm{dp}[x - 1][y + \mathrm{num}[x] - kx] \mathrel{+}= \mathrm{dp}[x][y] * \mathrm{choose} * \mathrm{arrange} $$ where $$ \mathrm{choose}(y + \mathrm{num}[x]人からkx人選ぶ): {}_{y + \mathrm{num}[x]} C_{kx} \\\ \mathrm{arrange}(kx人をx人kチームに分ける): \frac{(kx)!}{(x!)^k * k!} $$ で更新すればよい。
計算量もさっきと全く同じで$\mathcal{O}(n^2 \log(n))$。
実装
// mint, fact, fact_invは省略
signed main() {
int n; cin >> n;
init_factorial(n);
vector<int> num(n + 1);
REP (i, n) {
int x; cin >> x;
num[x]++;
}
vector<vector<mint>> dp(n + 1, vector<mint>(n + 1));
dp[n][0] = 1;
for (int x = n; x >= 1; x--) {
REP (y, n + 1) {
if (dp[x][y] == 0) continue;
for (int k = 0; y + num[x] - k * x >= 0; k++) {
mint choose = comb(y + num[x], k * x);
mint arrange = fact[k * x] * pow(fact_inv[x], k) * fact_inv[k];
dp[x - 1][y + num[x] - k * x] += dp[x][y] * choose * arrange;
}
}
}
cout << dp[0][0] << endl;
}
教訓
- DPはどの状態をまとめられるかが重要
- チーム分け(違うグループに属する人が存在すれば別カウント)は、$x$人のチームを作って$y$人残っている状態を同じとみなすとうまくいく
- 制約がきつい順に見ていけば、新しいものを見るときに以前の制約は全て無視できる