Rustコトハジメ

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

【yukicoder No.895 MESE】コンビネーションを注意深く計算する

No.895 MESE - yukicoder

問題

a+b+cビットの中に、aビット,bビット,cビット選び、それぞれx,y,zとする。x>y>zとなるようなzの総和を求めよ。

方針

f:id:akiradeveloper529:20190927233222j:plain

明らかにMSB(x)は最上位でないといけない。また、MSB(z)はd-2以下でないといけない。

まずMSB(z)を決めて、その下にc-1ビット選ぶ。これらの総和は、各ビットが何回出現するかを計算すればよい。MSBは、MSB(z)-1からc-1選ぶとおり回出現する。それ以下のビットは、そのビットが出現したと仮定して他はc-2ビットはどうでもいいと考えると、MSB(z)-2からc-2ビット選ぶとおり回出現することになる。

その上で、yのためにbビットとることを考える。ここでの制約は、MSB(z)より上位にMSB(y)を置くことだが、どこにMSB(y)を置くかを考える必要はなく、d-1-cからb個選ぶやり方から、MSB(y)>MSB(z)を満たさない組み合わせを除外すればよい。これは、MSB(z)より下にbビット全部を選んでしまう場合を計算すればよい。

実装

1-indexedの方がわかりよいと思った。

fn solve() {
    input!{
        a:u64,b:u64,c:u64,
    }
    let d = a+b+c;
    // d-2ビットからcビットとるやり方

    let mo = 1_000_000_007;
    let modcomb = ModComb::new(d, mo);
    let mut sum = 0;
    // 1-indexed
    // i: MSB of z
    for i in c..d-1 { // [c,d-2]
        let xmsbcnt = modcomb.nCk(i-1, c-1);
        let mut lsum = modpow(2,i-1,mo) * xmsbcnt;
        lsum %= mo;

        if c>=2 {
            // i以下のあるビットが出現する回数
            // ibit目とそいつを固定させて他からc-2個選ぶ
            let xothercnt = if c==2 {
                1
            } else {
                modcomb.nCk(i-2, c-2)
            };
            let sumother = (modpow(2,i-1,mo)+mo-1)%mo;
            lsum += sumother * xothercnt;
            lsum %= mo;
        }

        // yのとり方
        // d-1個からc個を引いたものからbを選ぶ
        // そこから、b<cとなってしまう場合を引く
        let ycnt = (modcomb.nCk(d-1-c, b) + mo - modcomb.nCk(i-c,b)) % mo;

        // 赤は自動的に決まる

        lsum *= ycnt;
        lsum %= mo;

        sum += lsum;
        sum %= mo;
    }
    println!("{}",sum);
}