Rustコトハジメ

プログラミング言語Rustに関する情報をお届けします。

包除原理を応用する問題

自力で解けず、他人の解法を見て、DPを使う解法はわかったのですが、理解しがたい解法があって、むしろそっちの方がマジョリティだったので、みんなが知ってる解法だから理解する必要があると考えて、結局5時間かかりました。コードから包除原理にノーヒントでたどり着けたことが逆に奇跡だと思います。

https://onlinejudge.u-aizu.ac.jp/problems/DPL_5_C

n個の区別できるボールをk個の区別できる箱に入れるとき、可能な入れ方の総数を求めてください。ただし、どの箱にも、1つ以上のボールを入れる。

という問題です。

ベースとなる考え方は、

  1. とにかくそれぞれのボールに1からkの数字をつけてみる (kn)
  2. この場合、空箱を含んでしまうから、空箱を含んでる場合を除く

です。ここで2を求めるのが仕事ということなります。

まず、DPを使った解法は、n個のボールをk-i個の箱で消費する小さな問題を解いて、それを大きな問題で使うということなので、これは容易に理解出来ます。

理解出来なかったのは、包除原理を使った解法です。

包除原理の意味と証明|思考力を鍛える数学

包除定理において、A(i)を「i番目の箱が空である」という事象とします。すると、左辺がまさに求めるものということになります。次に右辺において、Sum(|A(i)|)というのは、C(k,1) (k-1)nに等しいです。同様に、Sum(|A(i) かつ A(j)|)というのは、C(k,2) (k-2)nに等しいです。

これを愚直にコードにすると、DPなしで実装することが出来ます。

    input! {
        n: u64, k: u64
    }
    let p = 1_000_000_007;
    let mc = ModComb::new(1001,p);
    let res = modpow(k,n,p);
    let mut cnt = 0;
    for m in 1..k {
        let tmp = (mc.nCk(k, m) * modpow(k-m, n, p)) % p;
        if m % 2 == 1 {
            cnt = cnt + tmp;
        } else {
            cnt = (cnt + p - tmp) % p;
        }
    }
    let res = (res + p - cnt) % p;
    println!("{}", res);