母関数おじさんにとっては多分実家。
F - Paper Cutting
問題
長方形の中に縦線がH個、横線がW個ある。これからK本の辺を選んで、順番にその線に従って長方形をカットする。
各カットについて、そのスコアはカット後の分割された長方形の個数である。P(H + W, K)通りあるカットの方法について、スコアの合計の総和を求めよ。
解法
全てのカットが一様ランダムに選ばれるとした時の、スコアの合計の期待値を求めれば良い。
和の期待値は期待値の和なので、i回目のカットで縦線/横線を選んだときのスコアの期待値が分かれば良い。(iを0-indexedと思うことにして、0<=i<=K-1とする。)
i回目のカットが縦の場合を考える。i回目のカットの前に縦にj回カット、横に(i-j)回カットされていたとする。そのときの確率はC(H-1,j)*C(W,i-j)/C(H+W-1,i)で、スコアは(j+2)(i-j+1)である。これの積を全てのj (0 <= j <= i)について足し合わせれば良い。
ここで、係数がjの多項式なので、これを母関数を使って表すことにすると、C(H-1,j)(j+2)は(1/x) * (d/dx((1+x)^{h-1}x^2) = (1+x)^{h-2}((h-1)x+2(1+x))におけるx^jの係数、C(W,u)*(u+1)は(d/dx)((1+x)^w x) = (1+x)^{w-1}(wx + 1 + x)におけるx^uの係数である。求めたい和は畳み込みなので、結局これらの積(1+x)^{h+w-3}((h-1)x+2(1+x))wx + 1 + x)におけるx^iの係数である。これはあらかじめ階乗や階乗の逆数のテーブルを用意しておけば定数時間で計算できる。
i回目のカットが横の場合も同様である。i回目のカットが縦/横になる確率はそれぞれH/(H+W), W/(H+W)なので、これを掛けた値を合計すれば良い。
登場する典型
- 組み合わせの問題を母関数を使って表す
実装上の注意点
- H, Wが10^7とかなり大きいので、ループの中にO(log mod)の処理を乗せてはいけない!
- mod逆元など
- 定数時間でもmodは重いので、極力減らす!
提出:
Submission #4009374 - KEYENCE Programming Contest 2019 / キーエンス プログラミング コンテスト 2019
(Rust)
// ModInt省略 fn fact_init(w: usize) -> (Vec<ModInt>, Vec<ModInt>) { let mut fac = vec![ModInt::new(1); w]; let mut invfac = vec![0.into(); w]; for i in 1 .. w { fac[i] = fac[i - 1] * i as i64; } invfac[w - 1] = fac[w - 1].inv(); for i in (0 .. w - 1).rev() { invfac[i] = invfac[i + 1] * (i as i64 + 1); } (fac, invfac) } fn solve() { let out = std::io::stdout(); let mut out = BufWriter::new(out.lock()); macro_rules! puts { ($($format:tt)*) => (write!(out,$($format)*).unwrap()); } input! { h: usize, w: usize, k: usize, } let (fac, invfac) = fact_init(h + w + 1); let comb = |x, y| { if x < y { ModInt::new(0) } else { fac[x] * invfac[x - y] * invfac[y] } }; let invcomb = |x, y| { if x < y { ModInt::new(0) } else { invfac[x] * fac[x - y] * fac[y] } }; let mut tot = ModInt::new(0); let hp = ModInt::new(h as i64) * ModInt::new((h + w) as i64).inv(); let wp = ModInt::new(w as i64) * ModInt::new((h + w) as i64).inv(); for i in 0 .. k { let c0 = comb(h + w - 2, i) * 2; let c1 = if i > 0 { comb(h + w - 2, i - 1) } else { 0.into() }; let c2 = if i > 0 && h + w >= 3 { comb(h + w - 3, i - 1) } else { 0.into() }; let c3 = if i > 1 && h + w >= 3 { comb(h + w - 3, i - 2) } else { 0.into() }; let fac = invcomb(h + w - 1, i); let mut coef = c0; if i > 0 { coef += c1 * (w as i64 + 1) * 2; if h > 0 { coef += c2 * (h as i64 - 1); if i > 1 { coef += c3 * (h as i64 - 1) * (w as i64 + 1); } } } coef *= fac; tot += coef * hp; let mut coef = c0; if i > 0 { coef += c1 * (h as i64 + 1) * 2; if w > 0 { coef += c2 * (w as i64 - 1); if i > 1 { coef += c3 * (h as i64 + 1) * (w as i64 - 1); } } } coef *= fac; tot += coef * wp; } puts!("{}\n", tot * fac[h + w] * invfac[h + w - k]); }
まとめ
Eで1時間くらい使ってしまったせいで、これを40分で解けたのにコンテスト終了に間に合わせられなかった。