AtCoder Beginner Contest 404 E

https://atcoder.jp/contests/abc404/tasks/abc404_e

空でない茶碗は処理しないといけないので、空でない茶碗に移動します。問題は空の茶碗しかないときですが、そのときは空でない茶碗を右から探して、そこからDPで空でない茶碗にたどり着く最短手数を求めます。

// Bowls and Beans
#![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()
}


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

fn read_input() -> (Vec<usize>, Vec<i32>) {
    let _v: Vec<i32> = read_vec();
    let C: Vec<usize> = read_vec();
    let A: Vec<i32> = read_vec();
    (C, A)
}

fn find_non_zero(i: usize, C: &Vec<usize>, A: &Vec<i32>) -> usize {
    for j in (0..i-C[i]).rev() {
        if A[j] != 0 {
            return j
        }
    }
    return 0
}

use std::cmp::min;

// どこに豆を入れるか
fn find_pos(i: usize, C: &Vec<usize>, A: &Vec<i32>) -> (usize, u32) {
    if C[i] >= i {
        return (0, 1)
    }
    
    for j in (i-C[i]..i).rev() {
        if A[j] != 0 {
            return (j, 1)
        }
    }
    
    // 0しかないときが問題
    // 最後の0でない場所を探す
    let p = find_non_zero(i, C, A);
    let mut dp: Vec<u32> = vec![u32::MAX; i-p+1];
    dp[i-p] = 0;
    for j in (p+1..i+1).rev() { // jから移動する
        for k in 1..C[j]+1 {    // k個左に移動する
            if j < k || j - k < p {
                break
            }
            let q = j - k;      // qに移動
            dp[q-p] = min(dp[q-p], dp[j-p] + 1)
        }
    }
    (p, dp[0])
}

fn F(mut C: Vec<usize>, mut A: Vec<i32>) -> u32 {
    let N = C.len();
    C.insert(0, 0);
    A.insert(0, 0);
    let mut counter: u32 = 0;
    for i in (1..N+1).rev() {
        if A[i] == 0 {
            continue
        }
        let (j, n) = find_pos(i, &C, &A);
        A[j] += A[i];
        counter += n
    }
    counter
}

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