AtCoder Beginner Contest 378 E

https://atcoder.jp/contests/abc378/tasks/abc378_e

ふつうに計算すると計算量は、累積和を使ってもです。しかしよく考えると、一つのrについてで計算できます。
例えば、Mが5だとして、Aが[1, 2, 3, 3]だとして、rが3のとき、和は6, 5, 3だから、Mで割った余りは1, 0, 3です。rを4にすると、1, 0, 3に3-4をMで割った余りを追加して、そこに4を足せばよいです。しかし、Mで割った余りを足さなければならないので、

として、ここに4個分4を加えて3と4の2個だけMを引けばよいです。結局、r=4のときの和は、

となります。
すなわち、ある値より小さい値の個数と和が対数の計算量で得られる木を作ればよいです。

// Mod Sigma Problem
#![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()
}


//////////////////// Segment Tree ////////////////////

// (num, sum)
type Value = (usize, usize);

struct SegmentTree {
    n: usize,   // 2-pow
    v: Vec<Value>,
}

impl SegmentTree {
    fn add(&mut self, i: usize) {
        let mut k = i + self.n - 1;
        loop {
            let (n, s) = self.v[k];
            self.v[k] = (n + 1, s + i);
            if k == 0 {
                break
            }
            k = (k - 1) / 2;
        }
    }
    
    // v_0 + ... + v_{n-1}
    fn sum(&self, n: usize) -> Value {
        self.sum_core(n, 0, self.n, 0)
    }
    
    fn sum_core(&self, n: usize, first: usize, last: usize,
                                                i: usize) -> Value {
        let mid = (first + last) / 2;
        if n == last {
            self.v[i]
        }
        else if n <= mid {
            self.sum_core(n, first, mid, i*2+1)
        }
        else {
            let (n1, s1) = self.sum_core(mid, first, mid, i*2+1);
            let (n2, s2) = self.sum_core(n, mid, last, i*2+2);
            (n1+n2, s1+s2)
        }
    }
    
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn new(m: usize) -> SegmentTree {
        let n = SegmentTree::ceil_two_pow(m);
        let v: Vec<Value> = vec![(0, 0); n*2-1];
        SegmentTree { n, v }
    }
}


//////////////////// 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 F(M: usize, A: Vec<usize>) -> usize {
    let mut tree = SegmentTree::new(M);
    let mut counter: usize = 0;
    let mut acc: usize = 0;
    for i in 0..A.len() {
        let a = A[i] % M;
        acc = (acc + a) % M;
        let e = if a >= acc { a - acc } else { a + M - acc };
        tree.add(e);
        if acc == 0 {
            counter += tree.v[0].1
        }
        else {
            let (n, _) = tree.sum(M - acc);
            let (_, s1) = tree.sum(M);
            counter += s1 + (i + 1) * acc - M * (i + 1 - n)
        }
    }
    counter
}

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