AtCoder Beginner Contest 407 F

https://atcoder.jp/contests/abc407/tasks/abc407_f

これは典型的な和の順序を変えると解ける問題ですね。すなわち、ある幅を持ったウィンドウをスライドさせて最大値を調べるのではなく、ある値が最大値になる幅を調べます。入力例2で考えると、(値, 場所)を降順にソートして、まず6番目の5が最初なので、これが最大値になる幅の個数は、1が1個、2が2個、3が3個、4が3個、5が3個、6が3個、7が2個、8が1個となります。次に4番目の5で、5番目までしか範囲を広げられないので、1が1個、2が2個、3が2個、4が2個、5が1個となります。このように個数は線形の関数の3つ以下の組合せになることがわかります。そこで、線形の関数を持つ木を作って足していって、最後に答えを得ます。

// Sums of Sliding Window Maximum
#![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()
}


//////////////////// Value ////////////////////

use std::ops::Add;

// ax + b
#[derive(Clone, Copy)]
struct Value {
    a: i64,
    b: i64
}

impl Value {
    fn new(a: i64, b: i64) -> Value {
        Value { a, b }
    }
    
    fn apply(&self, x: usize) -> i64 {
        self.a * (x as i64) + self.b
    }
}

impl Add for Value {
    type Output = Self;
    
    fn add(self, other: Self) -> Self {
        Self { a: self.a + other.a, b: self.b + other.b }
    }
}

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

type Range = (usize, usize);

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

impl SegmentTree {
    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![Value::new(0, 0); n*2-1];
        SegmentTree { n, v }
    }
    
    fn is_leaf(&self, i: usize) -> bool {
        i >= self.n - 1
    }
    
    fn add_core(&mut self, f: Value, rng: Range,
                            i: usize, first: usize, last: usize) {
        let mid = (first + last) / 2;
        if self.is_leaf(i) {
            self.v[i] = self.v[i] + f
        }
        else if rng.0 <= first && last <= rng.1 {
            self.v[i] = self.v[i] + f
        }
        else {
            if rng.0 < mid {
                self.add_core(f, rng, i*2+1, first, mid)
            }
            if mid < rng.1 {
                self.add_core(f, rng, i*2+2, mid, last)
            }
        }
    }
    
    fn add(&mut self, f: Value, rng: Range) {
        self.add_core(f, rng, 0, 0, self.n)
    }
    
    fn get_value_core(&self, x: usize, i: usize,
                                first: usize, last: usize) -> i64 {
        let mid = (first + last) / 2;
        if self.is_leaf(i) {
            self.v[i].apply(x)
        }
        else if x < mid {
            self.v[i].apply(x) + self.get_value_core(x, i*2+1, first, mid)
        }
        else {
            self.v[i].apply(x) + self.get_value_core(x, i*2+2, mid, last)
        }
    }
    
    fn get_value(&self, x: usize) -> i64 {
        self.get_value_core(x, 0, 0, self.n)
    }
}


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

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

fn calc_distribution(m: usize, n: usize, a: i64) -> Vec<(Value, Range)> {
    if m > n {
        calc_distribution(n, m, a)
    }
    else if m == 0 {
        vec![(Value::new(0, a), (0, n+1))]
    }
    else if m + 1 >= n {
        vec![(Value::new(a, a), (0, m+1)),
             (Value::new(-a, a*((m+n+1) as i64)), (m+1, m+n+1))]
    }
    else {
        vec![(Value::new(a, a), (0, m+1)),
             (Value::new(0, a*((m+1) as i64)), (m+1, n+1)),
             (Value::new(-a, a*((m+n+1) as i64)), (n+1, m+n+1))]
    }
}

use std::collections::BTreeSet;

fn F(A: Vec<i64>) {
    let N = A.len();
    let mut B: Vec<(i64, usize)> = A.iter().cloned().zip(0..).collect();
    B.sort_by(|a, b| b.cmp(a));
    let mut tree = SegmentTree::new(N);
    let mut set: BTreeSet<usize> = BTreeSet::new();
    for (a, p) in B.into_iter() {
        let lower = set.range(..p).next_back();
        let upper = set.range(p..).next();
        let low = if lower.is_none() { 0 } else { lower.unwrap() + 1 };
        let upp = if upper.is_none() { N-1 } else { upper.unwrap() - 1 };
        let fs: Vec<(Value, Range)> = calc_distribution(p-low, upp-p, a);
        for (f, rng) in fs.into_iter() {
            tree.add(f, rng)
        }
        set.insert(p);
    }
    
    for x in 0..N {
        println!("{}", tree.get_value(x))
    }
}

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