AtCoder Beginner Contest 444 E

https://atcoder.jp/contests/abc444/tasks/abc444_e

左から見ていってそこをRとすると前回のLから右に動かします。そうしたときに、

を求めます。本当はバランス木を作ればいいのですが、面倒なのでAをソートしてSegment木を作って、その値が範囲にあるのかを個数で表します。もうちょっと簡単にできればよかったのですが。

// Sparse Range
#![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()
}


//////////////////// SegTree ////////////////////

trait Monoid {
    type S: Clone;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn e() -> Self::S;
}

struct SegTree<M: Monoid> {
    n: usize,
    size: usize,
    data: Vec<M::S>,
}

impl<M: Monoid> SegTree<M> {
    fn is_leaf(&self, i: usize) -> bool {
        i >= self.size - 1
    }
    
    fn new(n: usize) -> Self {
        let size = n.next_power_of_two();
        SegTree {
            n,
            size,
            data: vec![M::e(); size*2-1],
        }
    }
    
    fn from(v: Vec<M::S>) -> Self {
        let n = v.len();
        let size = n.next_power_of_two();
        let mut data = vec![M::e(); size*2-1];
        for i in 0..n {
            data[i+size-1] = v[i].clone();
        }
        let mut seg = SegTree { n, size, data };
        for i in (0..size-1).rev() {
            seg.data[i] = M::op(&seg.data[2*i+1], &seg.data[2*i+2]);
        }
        seg
    }
    
    fn set(&mut self, mut i: usize, x: M::S) {
        i = i + self.size - 1;
        self.data[i] = x;
        while i != 0 {
            i = (i - 1) / 2;
            self.data[i] = M::op(&self.data[2*i+1], &self.data[2*i+2]);
        }
    }
    
    fn prod(&self, mut l: usize, mut r: usize) -> M::S {
        let mut sml = M::e();
        let mut smr = M::e();
        l = l + self.size - 1;
        r = r + self.size - 1;
        while l < r {
            if l % 2 == 0 {
                sml = M::op(&sml, &self.data[l]);
                l += 1;
            }
            if r % 2 == 0 {
                r -= 1;
                smr = M::op(&self.data[r], &smr);
            }
            l = (l - 1) / 2;
            r = (r - 1) / 2
        }
        M::op(&sml, &smr)
    }
    
    fn get(&self, i: usize) -> M::S {
        self.data[i+self.size-1].clone()
    }
}


//////////////////// Monoid ////////////////////

use std::cmp::{min, max};

struct Max;
impl Monoid for Max {
    // ((min, minの元のindex), (max, maxの元のindex), 範囲にいくつ入っているか)
    type S = ((i32, usize), (i32, usize), usize);
    fn op((n1, x1, num1): &Self::S, (n2, x2, num2): &Self::S) -> Self::S {
        if *num1 > 0 && *num2 > 0 {
            (min(*n1, *n2), max(*x1, *x2), num1 + num2)
        }
        else if *num1 > 0 {
            (*n1, *x1, *num1)
        }
        else if *num2 > 0 {
            (*n2, *x2, *num2)
        }
        else {
            Max::e()
        }
    }
    fn e() -> Self::S { ((0, 0), (0, 0), 0) }
}


//////////////////// Tree ////////////////////

struct Tree {
    seg: SegTree::<Max>,
    ids: Vec<usize>,    // 元のindex -> SegTreeでのindex
}

impl Tree {
    fn set(&mut self, i: usize, b: bool) {
        let (n, _, _) = self.seg.get(self.ids[i]);
        let num = if b { 1 } else { 0 };
        self.seg.set(self.ids[i], (n, n, num))
    }
    
    fn left(&self, i: usize) -> Option<(i32, usize)> {
        let (_, x, num) = self.seg.prod(0, self.ids[i]);
        if num > 0 {
            Some(x)
        }
        else {
            None
        }
    }
    
    fn right(&self, i: usize) -> Option<(i32, usize)> {
        let (n, _, num) = self.seg.prod(self.ids[i]+1, self.ids.len());
        if num > 0 {
            Some(n)
        }
        else {
            None
        }
    }
    
    fn num_all(&self) -> usize {
        self.seg.data[0].2
    }
    
    fn create_ids(B: &Vec<(i32, usize)>) -> Vec<usize> {
        let mut ids: Vec<usize> = vec![0; B.len()];
        for (i, &(_, j)) in B.iter().enumerate() {
            ids[j] = i
        }
        ids
    }
    
    fn from(A: &Vec<i32>) -> Tree {
        let mut B: Vec<(i32, usize)> = A.iter().enumerate().
                                        map(|(i, &n)| (n, i)).collect();
        B.sort();
        let ids = Tree::create_ids(&B);
        let C: Vec<((i32, usize), (i32, usize), usize)> = B.into_iter().
                                map(|(n, i)| ((n, i), (n, i), 0)).collect();
        let seg = SegTree::from(C);
        Tree { seg, ids }
    }
}


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

fn read_input() -> (i32, Vec<i32>) {
    let v: Vec<i32> = read_vec();
    let D = v[1];
    let A: Vec<i32> = read_vec();
    (D, A)
}

fn F(D: i32, A: Vec<i32>) -> usize {
    let mut counter: usize = 0;
    let N = A.len();
    let mut tree = Tree::from(&A);
    let mut prev_j: usize = 0;
    for i in 0..N {
        tree.set(i, true);
        // 範囲内で最も近いものを探す
        let mut j: usize = 0;
        let mut j1 = i;
        loop {
            if let Some((a, new_j)) = tree.left(j1) {
                if a <= A[i] - D {
                    break
                }
                j1 = new_j;
                j = max(j, new_j)
            }
            else {
                break
            }
        }
        let mut j2 = i;
        loop {
            if let Some((a, new_j)) = tree.right(j2) {
                if a >= A[i] + D {
                    break
                }
                j2 = new_j;
                j = max(j, new_j)
            }
            else {
                break
            }
        }
        
        if j1 != i || j2 != i {
            for k in prev_j..j+1 {
                tree.set(k, false)
            }
            prev_j = j + 1
        }
        counter += tree.num_all()
    }
    counter
}

fn main() {
    let (D, A) = read_input();
    println!("{}", F(D, A))
}