D - ボールの箱詰め / Packing Balls into Boxes Editorial by admin
gemini-3.5-flash-thinking概要
区別できる \(K\) 個のボールを、制約条件を満たしつつ \(N\) 個の区別できる箱に「空箱がないように」分配する問題です。制約条件により、いくつかのボールは必ず同じ箱に入れる必要があります。
考察
1. 制約条件の整理(連結成分の抽出)
「ボール \(U_j\) とボール \(V_j\) は同じ箱に入れなければならない」という制約は、ボール同士のグループ化を意味しています。この関係は推移的(ボール \(A\) と \(B\)、ボール \(B\) と \(C\) が同じ箱に入るなら、ボール \(A\) と \(C\) も同じ箱に入る)であるため、ボールをいくつかのグループに分類することができます。
これは、ボールを頂点、制約を辺としたグラフにおける連結成分を求める問題、あるいはデータ構造の Union-Find (DSU) を用いてグループ分けする問題に帰着されます。
2. 問題の単純化
ボールをグループ分けした結果、全体で \(C\) 個のグループができたとします。 同じグループに属するボールは常に同じ行動をとる(同じ箱に入る)ため、各グループを「1つの大きなボール」とみなすことができます。
これにより、問題は以下のようにシンプルに言い換えることができます。 「区別できる \(C\) 個のグループを、区別できる \(N\) 個の箱に、空箱がないように分配する方法の総数を求めよ」
3. 空箱がないような分配方法(包除原理)
もし \(C < N\) であれば、グループの数が箱の数より少ないため、どのように分配しても必ず空箱ができてしまいます。したがって、この場合の答えは \(0\) です。
\(C \ge N\) の場合、空箱がないように分配する方法の総数を数える必要があります。これは包除原理を用いて効率的に計算できます。
空の箱の個数に着目します。 - すべての箱(\(N\) 個)に自由に分配する方法は \(N^C\) 通り。 - ここから、「少なくとも 1 つの箱が空になる」ような余分な割り当てを包除原理で排除します。
具体的には、空にする箱の個数を \(i\) としたとき、求める方法の総数は以下の式で表されます。 $\( \text{方法の総数} = \sum_{i=0}^{N} (-1)^i \binom{N}{i} (N-i)^C \)$
各項の意味は以下の通りです: - \(\binom{N}{i}\): 空にする \(i\) 個の箱の選び方 - \((N-i)^C\): 残りの \(N-i\) 個の箱に \(C\) 個のグループを自由に配置する方法(各グループの配置先が \(N-i\) 通りあるため、\((N-i)^C\) 通り) - \((-1)^i\): 包除原理による符号(\(i\) が奇数なら引き、偶数なら足す)
アルゴリズム
- グループ分け (Union-Find)
ボールを頂点とする Union-Find を用意します。与えられた \(M\) 個の制約 \((U_j, V_j)\) についてマージ操作(
merge)を行い、最終的な連結成分の個数 \(C\) を求めます。 - コーナーケースの処理
\(C < N\) の場合は、空箱をなくすことが不可能なため、
0を出力して終了します。 - 階乗テーブルの前計算
二項係数 \(\binom{N}{i}\) を高速に求めるために、階乗(
fact)と階乗の逆元(inv_fact)のテーブルを \(O(N)\) で前計算しておきます。 - 包除原理の計算
\(i = 0\) から \(N\) までループを回し、各ステップで \(\binom{N}{i} (N-i)^C\) を計算し、符号を考慮して足し引きします。\((N-i)^C\) の計算には、繰り返し二乗法(
pow)を用いることで高速に行います。
計算量
時間計算量: \(O(K + M \alpha(K) + N \log C)\)
- Union-Find のマージおよび連結成分数のカウントに \(O(K + M \alpha(K))\) かかります(\(\alpha\) は逆アッカーマン関数)。
- 階乗テーブルの構築に \(O(N)\) かかります。
- 包除原理のループは \(N\) 回回り、各ループで \((N-i)^C\) を \(O(\log C)\) で計算するため、ループ全体で \(O(N \log C)\) かかります。
- 全体として、制約 \(N, K, M \le 2 \times 10^5\) の下で実行時間制限に余裕で間に合います。
空間計算量: \(O(N + K)\)
- Union-Find の管理に \(O(K)\)、階乗テーブルの保持に \(O(N)\) のメモリを使用します。
実装のポイント
Union-Find (DSU) のライブラリ活用 AtCoder Library (ACL) の
atcoder::dsuを使用することで、連結成分の管理をバグなく簡潔に実装できます。自動的な mod 割り算(
modint) 計算途中で値が非常に大きくなるため、すべての計算をatcoder::modint998244353上で行うことで、オーバーフローを防ぎつつ自動的に \(998244353\) で割った余りを計算できます。ソースコード
#include <iostream>
#include <vector>
#include <atcoder/dsu>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
vector<mint> fact, inv_fact;
void init_fact(int n) {
fact.resize(n + 1);
inv_fact.resize(n + 1);
fact[0] = 1;
for (int i = 1; i <= n; i++) fact[i] = fact[i - 1] * i;
inv_fact[n] = fact[n].inv();
for (int i = n - 1; i >= 0; i--) inv_fact[i] = inv_fact[i + 1] * (i + 1);
}
mint nCr(int n, int r) {
if (r < 0 || r > n) return 0;
return fact[n] * inv_fact[r] * inv_fact[n - r];
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int N, K, M;
if (!(cin >> N >> K >> M)) return 0;
dsu d(K);
for (int i = 0; i < M; i++) {
int u, v;
cin >> u >> v;
u--; v--;
d.merge(u, v);
}
int C = 0;
for (int i = 0; i < K; i++) {
if (d.leader(i) == i) {
C++;
}
}
if (C < N) {
cout << 0 << "\n";
return 0;
}
init_fact(N);
mint ans = 0;
for (int i = 0; i <= N; i++) {
mint term = nCr(N, i) * mint(N - i).pow(C);
if (i % 2 == 1) {
ans -= term;
} else {
ans += term;
}
}
cout << ans.val() << "\n";
return 0;
}
この解説は gemini-3.5-flash-thinking によって生成されました。
posted:
last update: