AtCoder Beginner Contest 388 D

https://atcoder.jp/contests/abc388/tasks/abc388_d

石の数を変化させる順番を変えると組みやすくなります。
入力例1は、

5 0 9 3

から、最初の宇宙人は最終的に3つ石を2～4番目の宇宙人に配るので、

2 1 10 4

となります。2番目の宇宙人は2つ石を配りたいところですが、1個しかないので3番目の宇宙人に配るので、

2 0 11 4

となります。

なので、ある宇宙人が持つ石の数を得る、ある宇宙人の石の数をセットする、ある範囲の宇宙人の石の数を一つずつ増やす、という3つメソッドを実装した木を作ればよいです。
Segment Treeを書けばよさそうなのですが、なぜかTLEになってしまいます。手元だと全然そんなことないのですが。よく考えると、宇宙人が石を配る段階でその宇宙人の最終的な石の数が確定するので、最後に石の数を得る必要は無いです。そうすると、インクリメントする範囲は最初からにすると木の更新が速くなります。これで余裕で通るようになりました。

// Coming of Age Celebration
#![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 ////////////////////

type Value = i32;
type Range = (usize, usize);

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

impl SegmentTree {
    fn get(&self, i: usize) -> Value {
        let mut value: Value = 0;
        let mut k: usize = 0;
        let mut first: usize = 0;
        let mut last: usize = self.n;
        while k < self.n - 1 {
            value += self.v[k];
            let mid = (first + last) / 2;
            if i < mid {
                k = k * 2 + 1;
                last = mid
            }
            else {
                k = k * 2 + 2;
                first = mid
            }
        }
        value + self.v[k]
    }
    
    fn set(&mut self, i: usize, mut value: Value) {
        let mut k: usize = 0;
        let mut first: usize = 0;
        let mut last: usize = self.n;
        while k < self.n - 1 {
            value -= self.v[k];
            let mid = (first + last) / 2;
            if i < mid {
                k = k * 2 + 1;
                last = mid
            }
            else {
                k = k * 2 + 2;
                first = mid
            }
        }
        self.v[i+self.n-1] = value
    }
    
    fn increment(&mut self, rng: Range) {
        if rng == (0, self.m) {
            self.v[0] += 1
        }
        else {
            self.increment_core(rng, (0, self.n), 0)
        }
    }
    
    fn increment_core(&mut self, rng: Range, tree_rng: Range, i: usize) {
        let (first, last) = tree_rng;
        if rng.0 <= first && last <= rng.1 {
            self.v[i] += 1
        }
        else {
            let mid = (first + last) / 2;
            if rng.1 <= mid {
                self.increment_core(rng, (first, mid), i*2+1)
            }
            else if rng.0 >= mid {
                self.increment_core(rng, (mid, last), i*2+2)
            }
            else {
                self.increment_core(rng, (first, mid), i*2+1);
                self.increment_core(rng, (mid, last), i*2+2)
            }
        }
    }
    
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn create(a: &Vec<Value>) -> SegmentTree {
        let m = a.len();
        let n = SegmentTree::ceil_two_pow(m);
        let mut v: Vec<Value> = vec![0; n*2-1];
        for i in n-1..n+m-1 {
            v[i] = a[i+1-n]
        }
        SegmentTree { n, m, v }
    }
}


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

fn read_input() -> Vec<i32> {
    let _N: usize = read();
    let A: Vec<i32> = read_vec();
    A
}

fn F(A: Vec<i32>) {
    let N = A.len();
    let mut tree = SegmentTree::create(&A);
    for i in 0..N-1 {
        let n = tree.get(i) as usize;
        let r = N - i - 1;
        if n >= r {
            print!("{} ", n - r);
            tree.increment((0, tree.n));
            tree.set(i, (n - r) as Value)
        }
        else {
            print!("{} ", 0);
            tree.increment((0, i + 1 + n));
            tree.set(i, 0)
        }
    }
    println!("{}", tree.get(N-1))
}

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