https://atcoder.jp/contests/abc359/tasks/abc359_e
左で一番最初に自分以上の値を探して、そこまでを自分の値と同じにします。そういう木を作ればよいです。
// Water Tank #![allow(non_snake_case)] use std::cmp::{min, max}; //////////////////// 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 = i64; struct SegmentTree { n: usize, // 2-pow m: usize, mins: Vec<Value>, maxs: Vec<Value>, sums: Vec<Value>, } impl SegmentTree { // [0, j)のeより小さいところをeにする fn set_max(&mut self, j: usize, e: Value) { self.set_max_core(j, e, 0, 0, self.n); } fn set_max_core(&mut self, j: usize, e: Value, cell: usize, first: usize, last: usize) { if j <= first { } else if last - first == 1 && self.mins[cell] < e { self.mins[cell] = e; self.maxs[cell] = e; self.sums[cell] = e; } else if self.maxs[cell] <= e && last <= j { self.mins[cell] = e; self.maxs[cell] = e; self.sums[cell] = e * ((last - first) as Value) } else if self.mins[cell] < e { let mid = (first + last) / 2; self.set_max_core(j, e, cell*2+1, first, mid); self.set_max_core(j, e, cell*2+2, mid, last); self.mins[cell] = min(self.mins[cell*2+1], self.mins[cell*2+2]); self.maxs[cell] = max(self.maxs[cell*2+1], self.maxs[cell*2+2]); self.sums[cell] = self.sums[cell*2+1] + self.sums[cell*2+2] } } // [0, j)での和 fn sum(&self) -> Value { self.sums[0] } fn ceil_two_pow(n: usize) -> usize { if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 } } fn create(m: usize) -> SegmentTree { let n = SegmentTree::ceil_two_pow(m); let mins: Vec<Value> = vec![0; n*2-1]; let maxs: Vec<Value> = vec![0; n*2-1]; let sums: Vec<Value> = vec![0; n*2-1]; SegmentTree { n, m, mins, maxs, sums } } } //////////////////// process //////////////////// fn read_input() -> Vec<Value> { let _N: usize = read(); let H: Vec<Value> = read_vec(); H } fn print_vec(v: Vec<Value>) { let N = v.len(); for i in 0..N-1 { print!("{} ", v[i]) } println!("{}", v[N-1]) } fn F(H: Vec<Value>) { let N = H.len(); let mut tree = SegmentTree::create(N); let mut v: Vec<Value> = vec![0; N]; for (i, h) in H.into_iter().enumerate() { tree.set_max(i+1, h); v[i] = tree.sum() + 1 } print_vec(v) } fn main() { let H = read_input(); F(H) }