AtCoder Beginner Contest 340 E

https://atcoder.jp/contests/abc340/tasks/abc340_e

入力例1で、初回で2 2 0 5 6になって、2回目は6を取り出して、0に2を1～4に1を足します。このように範囲に同じ値を足すので、そういう木を作ればよいです。

// Mancala 2
#![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 ////////////////////

struct SegmentTree {
    n: usize,
    m: usize,
    v: Vec<i64>,
}

impl SegmentTree {
    fn add_range(&mut self, a: i64, first: usize, last: usize) {
        if first < last {
            self.add_range_core(a, (first, last), 0, 0, self.n)
        }
    }
    
    fn add_range_core(&mut self, a: i64, range: (usize, usize),
                            i: usize, first: usize, last: usize) {
        let (f, l) = range;
        if f <= first && last <= l {
            self.v[i] += a
        }
        else {
            let mid = (first + last) / 2;
            if f < mid {
                self.add_range_core(a, range, i*2+1, first, mid)
            }
            if mid < l {
                self.add_range_core(a, range, i*2+2, mid, last)
            }
        }
    }
    
    fn get_core(&self, i: usize) -> i64 {
        if i == 0 {
            self.v[i]
        }
        else {
            self.v[i] + self.get_core((i-1)/2)
        }
    }
    
    fn get(&self, i: usize) -> i64 {
        self.get_core(i + self.n - 1)
    }
    
    fn set(&mut self, a: i64, i: usize) {
        let mut j = i + self.n - 1;
        let mut s = 0;
        loop {
            s += self.v[j];
            if j == 0 {
                break
            }
            j = (j-1) / 2;
        }
        self.v[i+self.n-1] += a - s
    }
    
    fn print(&self) {
        println!("{}", (0..self.m).map(|i| self.get(i).to_string()).
                                        collect::<Vec<_>>().join(" "))
    }
    
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn create(A: Vec<i64>) -> SegmentTree {
        let m = A.len();
        let n = SegmentTree::ceil_two_pow(m);
        let mut v: Vec<i64> = vec![0; n*2-1];
        for i in 0..m {
            v[i+n-1] = A[i]
        }
        SegmentTree { n, m, v }
    }
}


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

fn read_input() -> (Vec<i64>, Vec<usize>) {
    let _v: Vec<usize> = read_vec();
    let A = read_vec();
    let B = read_vec();
    (A, B)
}

fn F(A: Vec<i64>, B: Vec<usize>) {
    let N = A.len();
    let mut tree = SegmentTree::create(A);
    for b in B.into_iter() {
        let C = tree.get(b);
        tree.set(0, b);
        let i = (b + 1) % N;
        let q = C / (N as i64) ;
        let r = (C as usize) % N;
        if i + r <= N {
            tree.add_range(q, 0, i);
            tree.add_range(q+1, i, i+r);
            tree.add_range(q, i+r, N)
        }
        else {
            tree.add_range(q+1, 0, i+r-N);
            tree.add_range(q, i+r-N, i);
            tree.add_range(q+1, i, N)
        }
    }
    tree.print()
}

fn main() {
    let (A, B) = read_input();
    F(A, B)
}