Official

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\) が奇数なら引き、偶数なら足す)

アルゴリズム

  1. グループ分け (Union-Find) ボールを頂点とする Union-Find を用意します。与えられた \(M\) 個の制約 \((U_j, V_j)\) についてマージ操作(merge)を行い、最終的な連結成分の個数 \(C\) を求めます。
  2. コーナーケースの処理 \(C < N\) の場合は、空箱をなくすことが不可能なため、 0 を出力して終了します。
  3. 階乗テーブルの前計算 二項係数 \(\binom{N}{i}\) を高速に求めるために、階乗(fact)と階乗の逆元(inv_fact)のテーブルを \(O(N)\) で前計算しておきます。
  4. 包除原理の計算 \(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: