AtCoder Beginner Contest 339 E

https://atcoder.jp/contests/abc339/tasks/abc339_e

簡単なDPです。まで最大の列の長さをとすると、

となります。ただ、このままだと計算量がで間に合わないので、範囲の最大値を取れる木を用意します。

// Smooth Subsequence
#![allow(non_snake_case)]

use std::cmp::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 ////////////////////

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

impl SegmentTree {
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn max_range(&self, range: (usize, usize)) -> i32 {
        self.max_range_core(range, 0, 1, self.n+1)
    }
    
    fn max_range_core(&self, range: (usize, usize),
                        i: usize, first: usize, last: usize) -> i32 {
        let (f, l) = range;
        if f <= first && last <= l {
            self.v[i]
        }
        else {
            let mid = (first + last) / 2;
            let m1 = if f < mid {
                        self.max_range_core(range, i*2+1, first, mid)
                    } else { 0 };
            let m2 = if mid < l {
                        self.max_range_core(range, i*2+2, mid, last)
                    } else { 0 };
            max(m1, m2)
        }
    }
    
    // set v to i
    fn set(&mut self, i: usize, v: i32) {
        let mut j = i + self.n - 1;
        if j >= self.n {
            self.v[j-1] = v
        }
        while j > 1 {
            j >>= 1;
            self.v[j-1] = max(self.v[j-1], v)
        }
    }
    
    fn create(m: usize) -> SegmentTree {
        let n = SegmentTree::ceil_two_pow(m);
        let v: Vec<i32> = vec![0; n*2-1];
        SegmentTree { n, v }
    }
}


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

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

fn get_range(a: usize, D: usize, M: usize) -> (usize, usize) {
    let first = if a < D + 1 { 1 } else { a - D };
    let last = if a + D > M { M + 1 } else { a + D + 1 };
    (first, last)
}

fn F(D: usize, A: Vec<usize>) -> i32 {
    let M = *A.iter().max().unwrap();
    let mut tree = SegmentTree::create(M);
    for a in A.into_iter() {
        let s = tree.max_range(get_range(a, D, M));
        tree.set(a, s+1)
    }
    tree.max_range((1, M+1))
}

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