模擬国内2017G へやわり!(Room Assignment)

前置き

解説の天才二項係数が理解できなかったので,形式的べき級数を用いた機械的(?)な解法の導出について書きます.

問題

onlinejudge.u-aizu.ac.jp

解法

$(i, a _ i)$ の辺を張った無向グラフ $G$ を考える.$G$ から多重辺を取り除いてできる新しいグラフ $G'$ に関して,$G'$ がいくつかのパスの集合となっていなければ valid な部屋割りは存在しないので,答えは $0$ である.以下,$G'$ がパスの集合である場合を考える.

条件「任意の二つの異なる提案された部屋割り A, B に対して,部屋割り A の空き部屋に誰か1人を移しても部屋割り B にならない」について,長さ $3$ 以上のパスに属する人を移動させることで他の valid な部屋割りを得られることはないので,長さ $2$ のパスに属する人を移動させる場合のみを考えればよい.$a _ i \neq i$ より, 長さ $1$ のパスは存在しないことに注意する.

長さ $2$ のパスの個数を $N$, 長さ $3$ 以上のパスの個数を $M $ とおく.$N + M $ 個のパスたちを並べた後に空き部屋を挿入することで部屋割りを得ることを考える.一旦パスの向きを無視して,最後に $2 ^ {N + M}$ を答えに掛けることにする.

例えば [oo][oo] ([oo] は長さ $2$ のパスを表す) に対しては,空き部屋の挿入方法は x[oo][oo], [oo]x[oo], [oo][oo]x の $3$ つある.x[oo][oo][oo]x[oo][oo]x[oo][oo][oo]x は人を $1$ 人選んで動かすことで得られる組合せなので,同時に部屋割りの集合に含めることは出来ない.従って,条件に違反しないようにできるだけ多くの部屋割りを選ぶには,x[oo][oo], [oo][oo]x の $2$ つを選ぶことになる.一般に,[oo] が $i$ 個続く場合,空き部屋の挿入方法は最大で $i + 1 - \lceil i / 2 \rceil$ 個選択することが出来ることが分かる.

さらに,例えば [ooo][oo][oo][ooo][ooo][oo][oo][oo] に空き部屋を挿入することを考えると,[oo] が連続する極大な区間たちに注目することで,この場合に空き部屋の挿入の仕方は最大で $N + M + 1 - \lceil 2 / 2 \rceil - \lceil 3 / 2 \rceil$ 個選択することが出来ることが分かる.

以上の観察を元に,$M $ 個の長さ $3$ 以上のパスを予め並べておき,$M+1$ 個の隙間に合計 $N$ 個の長さ $2$ のパスを挿入することを考える.それぞれの隙間に長さ $2$ のパスが幾つ挿入されるかのみに注目すればよいので,パスが選ばれる順番は区別せずに数えた後に $N! M!$ を答えに掛けることにする.

動的計画法により数えることを考える.

$$\begin{aligned} \mathrm{dp}[i][j][k] = &\text{$i$ 個の隙間に合計 $j$ 個の長さ $2$ のパスを挿入する方法であって,}\\ &\text{空き部屋の挿入方法を最大 $i + j - k$ 個選べるようなものの個数} \end{aligned}$$

とすると,次のように初期化および遷移を書くことが出来る.求めたいのは $S := \displaystyle \sum _ {k} (N+M+1-k)\mathrm{dp}[M+1][N][k]$ である.

$$ \mathrm{dp}[i][j][k] = \begin{cases} 1 & \text{if } i = 0\land j = 0\land k = 0 \\ 0 & \text{if } i = 0\land (j \neq 0 \lor k \neq 0)\\ \displaystyle \sum _ {p = 0} ^ {\infty} \mathrm{dp}[i - 1][j - p][k - \lceil p / 2 \rceil] & \text{otherwise} \end{cases}. $$

従って,形式的べき級数 $\displaystyle f _ {i}(x,y) = \sum _ j \sum _ k \mathrm{dp}[i][j][k] x ^ j y ^ k$ を定義すると,以下が成り立つ.

$$\begin{aligned} f _ 0 &= 1, \\ f _ {i + 1} &= f _ i \cdot (x ^ 0 y ^ 0 + x ^ 1 y ^ 1 + x ^ 2 y ^ 1 + x ^ 3 y ^ 2 + x ^ 4 y ^ 2 + \cdots) \\ &= f _ i \left(\sum _ {i = 0} ^ {\infty} x ^ {2i} y ^ i + xy\sum _ {i = 0} ^ {\infty} x ^ {2i} y ^ i\right) \\ &= f _ i \cdot \dfrac{1 + xy}{1 - x ^ 2 y}. \end{aligned}$$

即ち,$f _ {M + 1} = \left(\dfrac{1 + xy}{1 - x ^ 2 y}\right) ^ {M + 1}$ を得る.

ここで,$S = \displaystyle \sum _ {k} (N+M+1-k) [x ^ N y ^ k] f _ {M + 1}$ が畳込みの形をしていることに注目する.$\displaystyle \sum _ {i} (i + 1) y ^ i = \dfrac{1}{(1 - y) ^ 2}$ であるから,次が成り立つ.

$$ S = [x ^ N y ^ {N + M}] \dfrac{1}{(1 - y) ^ 2}\cdot \left(\dfrac{1 + xy}{1 - x ^ 2 y}\right) ^ {M + 1}. $$

右辺の $x ^ N$ の係数に注目する.

$$\begin{aligned} (1 + xy) ^ {M + 1} &= \sum _ {i = 0} ^ {M + 1}\binom{M + 1}{i} y ^ i \cdot x ^ i, \\ \dfrac{1}{(1 - x ^ 2 y) ^ {M + 1}} &= \sum _ {i = 0} ^ \infty \binom{M + i}{i} y ^ i \cdot x ^ {2i}, \end{aligned}$$

であるから,次を得る.

$$\begin{aligned} S = [y ^ {N + M}] \dfrac{1}{(1 - y) ^ 2} \sum _ {i} \binom{M + 1}{N - 2i} \binom{M + i}{i} y ^ {N - i} \end{aligned}$$

階乗を前計算しておくことで,これは $\Theta(N + M)$ 時間で計算できる.求める答えは $2 ^ {N + M} N! M! S$ であるから,全体 $\Theta(N + M + \log \mathrm{mod})$ で本問題を解くことが出来た.

実装

注: AtCoder Library 使用

#include <iostream>

#include <atcoder/modint>
#include <atcoder/dsu>

using mint = atcoder::modint1000000007;

void remove_multiedges(std::vector<std::vector<int>>& g) {
    std::vector<uint8_t> exists(g.size(), 0);
    for (auto& vs : g) {
        for (int v : vs) exists[v] = true;
        vs.erase(std::remove_if(vs.begin(), vs.end(), [&](int v) { return not std::exchange(exists[v], false); }), vs.end());
    }
}

constexpr int L = 200000;

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    std::vector<mint> fac(L + 1), fac_inv(L + 1);
    fac[0] = 1;
    for (int i = 1; i <= L; ++i) {
        fac[i] = fac[i - 1] * i;
    }

    fac_inv[L] = fac[L].inv();
    for (int i = L; i > 0; --i) {
        fac_inv[i - 1] = fac_inv[i] * i;
    }

    auto binom = [&fac, &fac_inv](int n, int r) -> mint {
        if (r < 0 or r > n) return 0;
        return fac[n] * fac_inv[r] * fac_inv[n - r];
    };
    
    while (true) {
        int k;
        std::cin >> k;

        if (k == 0) break;

        std::vector<int> a(k);
        for (auto& e : a) std::cin >> e, --e;

        atcoder::dsu uf(k);

        std::vector<std::vector<int>> g(k);
        for (int i = 0; i < k; ++i) {
            g[i].push_back(a[i]);
            g[a[i]].push_back(i);
            uf.merge(i, a[i]);
        }
        remove_multiedges(g);

        int n = 0, m = 0;

        bool all_path = true;
        for (auto group : uf.groups()) {
            n += group.size() == 2;
            m += group.size() >= 3;

            int c1 = 0, c2 = 0;
            for (int v : group) {
                c1 += g[v].size() == 1;
                c2 += g[v].size() == 2;
                all_path &= g[v].size() <= 2;
            }
            all_path &= c1 == 2;
        }

        if (all_path) {
            std::vector<mint> f(n + m + 1);
            for (int i = 0; 2 * i <= n; ++i) {
                f[n - i] = binom(m + i, m) * binom(m + 1, n - 2 * i);
            }

            std::partial_sum(f.begin(), f.end(), f.begin());
            std::partial_sum(f.begin(), f.end(), f.begin());

            std::cout << (mint(2).pow(n + m) * fac[n] * fac[m] * f[n + m]).val() << '\n';
        } else {
            std::cout << 0 << '\n';
        }
    }
    return 0;
}