KEYENCE Programming Contest 2019 F - Paper Cutting

母関数おじさんにとっては多分実家。
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分で解けたのにコンテスト終了に間に合わせられなかった。