AtCoder Beginner Contest 367 D

https://atcoder.jp/contests/abc367/tasks/abc367_d

累積和を計算して、Mの剰余が同じ位置を集めます。そしてその位置の集まりで差がN以下の個数をカウントします。そのとき、その集まりが大きくなる可能性がありますが、尺取り法的に計算すれば速くなります。

// Pedometer
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let M = v[1];
    let A: Vec<usize> = read_vec();
    (M, A)
}

fn accumulate(A: Vec<usize>, M: usize) -> Vec<Vec<usize>> {
    let mut acc: Vec<Vec<usize>> = vec![vec![]; M*2+1];
    let mut r: usize = 0;
    acc[0].push(0);
    for (i, a) in A.iter().enumerate() {
        r = (r + a) % M;
        acc[r].push(i+1)
    }
    for (i, a) in A.iter().enumerate() {
        r = (r + a) % M;
        acc[r].push(i+A.len()+1)
    }
    acc
}

fn F_each(ps: &Vec<usize>, N: usize) -> usize {
    let mut counter = 0;
    let mut k: usize = 0;
    let L = ps.len();
    for (i, &p) in ps.iter().enumerate() {
        if p >= N {
            break
        }
        while k < L && p + N > ps[k] {
            k += 1
        }
        counter += k - i - 1
    }
    counter
}

fn F(M: usize, A: Vec<usize>) -> usize {
    let N = A.len();
    let acc = accumulate(A, M);
    let mut counter: usize = 0;
    for i in 0..M {
        counter += F_each(&acc[i], N)
    }
    counter
}

fn main() {
    let (M, A) = read_input();
    println!("{}", F(M, A))
}