Rustコトハジメ

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

【ABC105-D Candy Distribution】累積和とmod

見るからに典型問題だけど、試行錯誤して解答を思いつくのに30分くらいかかったのでメモ。

D - Candy Distribution

問題

長さN(105)の配列A(要素は109)と整数M(109)が与えられる。[l,r]の和がMの倍数になる(l,r)の総数を求めよ。

方針

二乗は明らかに無理。とりあえずl,rの計算をするには累積和を計算すれば出来ることはわかる。なんとかここを手がかりにしたい。

累積和からl,rを求めた時にmod Mということは、[0,l-1]と[0,r]がmod Mで等しいということなので、あとはHashMapを持って舐めていけば終わる(要素数はたかだか105なのでMLEしない。配列で持つとMLE)。

実装

「ない」ことをインデックス0で表すためにiまでの和をi+1に記録するようにしている。

fn solve() {
    input!{
        N:usize,M:usize,
        A:[usize;N],
    }
    let mut cumsum = vec![0;N+1];
    for i in 0..N{
        let a=A[i];
        cumsum[i+1]=(cumsum[i]+a)%M;
    }
    // dbg!(&cumsum);
    let mut hs = HashMap::new();
    let mut res: usize = 0;
    for i in 1..N+1 {
        let m = cumsum[i];
        // 自分自身がMの倍数だったらカウント
        if m == 0 { res += 1; }
        // あとは自分より前に自分と同じmod Mのやつを探す
        res += *hs.get(&m).unwrap_or(&0);
        *hs.entry(m).or_insert(0) += 1;
    }
    println!("{}",res);
}

考察

なんとかの倍数系問題は、差分に注目することが多いような気がする。今回の問題でいうと、mod Mのカウントまで情報を圧縮してしまえるから、高速化が出来る。

あとは、足し算系なら、

D - Five, Five Everywhere

のようなタイプとか。