Mod Sigma Problem の解説

解説を書くほど言語化できていない。

左からなめていって、現在のbaseを持ちつつ和をセグ木で取得する。

fn main() {
    input! {
        n: usize,
        m: usize,
        a: [usize; n]
    }
    let acc = {
        let a = a.iter().map(|&e| e % m).collect_vec();
        let mut acc = vec![0];
        for &num in a.iter() {
            acc.push((acc.last().unwrap() + num) % m);
        }
        acc.remove(0);
        assert!(acc.len() == n);
        acc
    };
    
    let mut seg = Segtree::<M>::new(m);
    let mut seg2 = Segtree::<M>::new(m);
    for &num in acc.iter() {
        let pre = seg.get(num);
        seg.set(num, pre + 1);
        let pre2 = seg2.get(num);
        seg2.set(num, pre2 + num);
    }
    let mut base = 0;
    let mut ans = 0;
    for (idx, &num) in a.iter().enumerate() {
        let upcount = seg.prod(base..);
        let upsum = seg2.prod(base..);
        let downcount = seg.prod(..base);
        let downsum = seg2.prod(..base);


        ans += upsum - upcount * base;
        let pos = downsum + downcount * m;
        ans += pos - downcount * base;


        base += num;
        base %= m;

        let p1 = seg.get(acc[idx]);
        seg.set(acc[idx], p1 - 1);
        let p2 = seg2.get(acc[idx]);
        seg2.set(acc[idx], p2 - acc[idx]);
    }
    println!("{}", ans);
}