AtCoder Beginner Contest 359 E

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)
}