AtCoder Beginner Contest 334 F

https://atcoder.jp/contests/abc334/tasks/abc334_f

DPですが、状態が今いる位置と持っているプレゼント数なので、状態数がO(NK)なので間に合いません。しかし、スタートに戻る以外は次に移るので同じ値を足すだけです。また、スタートに戻るのは各状態から可能なので、最小値を取る必要があります。なので、セグメントツリーを実装すればよいです。あと辿っているとプレゼント数が一つずつ減っていくので、ずれていくのがめんどうですね。

// Christmas Present 2
#![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()
}


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

const INF: f64 = 1e18;

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

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<f64> = vec![INF; n*2-1];
        SegmentTree { n, v, d: 0.0 }
    }
    
    fn min(a: f64, b: f64) -> f64 {
        if a < b { a } else { b }
    }
    
    fn add_all(&mut self, a: f64) {
        self.d += a
    }
    
    fn update(&mut self, k: usize, a: f64) {
        let mut i = k + self.n - 1;
        self.v[i] = a - self.d;
        while i > 0 {
            i = (i - 1) / 2;
            self.v[i] = SegmentTree::min(self.v[i*2+1], self.v[i*2+2])
        }
    }
    
    fn range_min(&self, i: usize, j: usize) -> f64 {
        self.query(i, j, 0, 0, self.n) + self.d
    }
    
    fn query(&self, i: usize, j: usize, k: usize, l: usize, r: usize) -> f64 {
        if r <= i || j <= l {
            INF
        }
        else if i <= l && r <= j {
            self.v[k]
        }
        else {
            let mid = (l + r) / 2;
            let m1 = self.query(i, j, k*2+1, l, mid);
            let m2 = self.query(i, j, k*2+2, mid, r);
            SegmentTree::min(m1, m2)
        }
    }
}


//////////////////// Point ////////////////////

type Point = (f64, f64);

fn read_point() -> Point {
    let v = read_vec::<f64>();
    (v[0], v[1])
}

fn dist(pt1: Point, pt2: Point) -> f64 {
    ((pt1.0 - pt2.0).powf(2.0) + (pt1.1 - pt2.1).powf(2.0)).sqrt()
}


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

fn read_input() -> (usize, Vec<Point>) {
    let v = read_vec::<usize>();
    let (N, K) = (v[0], v[1]);
    let points: Vec<Point> = (0..N+1).map(|_| read_point()).collect();
    (K, points)
}

fn initialize_dp(K: usize) -> SegmentTree {
    let mut dp = SegmentTree::new(K);
    dp.update(K-1, 0.0);
    dp
}

fn update_dp(dp: &mut SegmentTree, i: usize, points: &Vec<Point>, K: usize) {
    let d = dist(points[i-1], points[i]);
    dp.add_all(d);
    let d1 = dist(points[i-1], points[0]);
    let d2 = dist(points[0], points[i]);
    let p = (K - i % K) % K;    // tail
    dp.update(p, dp.range_min(0, K) + d1 + d2 - d);
}

fn F(K: usize, points: Vec<Point>) -> f64 {
    let N = points.len();
    let mut dp = initialize_dp(K);
    for i in 1..N {
        update_dp(&mut dp, i, &points, K);
    }
    let d = dist(points[N-1], points[0]);
    dp.range_min(0, K) + d
}

fn main() {
    let (K, points) = read_input();
    println!("{}", F(K, points))
}