Rustコトハジメ

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

【DISCO presents ディスカバリーチャンネル コードコンテスト2016 予選 C ロト2】約数の数に問題を圧縮する系

なるほどなーと思う問題でした。自力AC出来なかったです。109の約数がたかだか1300個程度しかないという知識を活かして問題を圧縮する系です。今回もそんな臭いはしたけど、一歩踏み込めなかったです。

C - ロト2

問題

長さNの数列aが与えられる。このうちaiajがKの倍数となるようなi,jの組み合わせは何通りあるか?

方針

gcd(ai,K)gcd(aj,K)がKの倍数 => aiajがKの倍数

に気づけばよいはず。

gcdというのは、素因数分解したその上でビットマスクするようなものですから、gcd(ai,K)とgcd(aj,K)でKの分の素因数を全部かき集められるなら、aiajがKの倍数になってそうです。

Kの約数は1300個程度しかないため、aiをgcd(ai,K)によって分類することが出来ます。

あとはグループを組み合わせるだけ。

gcd(ai,K)は一意に定まるため、同じiが複数のグループに属しないというのもポイントです。

実装

最初resがi32で推論されていてWAしました。

返す値は明示的にusizeなりu64なり型を指定する方が良いと思います。その他のコードで別の型に推論されているなら、コンパイル時にわかるため、出口で指定してしまうのが一番確実と思います。

fn solve() {
    input!{
        N:usize,K:usize,
        a:[usize;N],
    }
    let mut cnt = HashMap::new();
    for ai in a {
        *cnt.entry(gcd(ai,K)).or_insert(0 as usize) += 1;
    }
    // dbg!(&cnt);
    let mut res: usize = 0;
    for (k0,v0) in &cnt {
        for (k1,v1) in &cnt {
            if k0>k1 { continue; }
            if (k0*k1)%K==0 {
                if k0==k1 {
                    res+=v0*(v0-1)/2;
                } else {
                    res+=v0*v1;
                }
            }
        }
    }
    println!("{}",res);
}