AtCoder Beginner Contest 323 E

https://atcoder.jp/contests/abc323/tasks/abc323_e

ある時刻tに曲が終わる確率をDPで求めればよいです。

// Playlist
#![allow(non_snake_case)]


//////////////////// constances ////////////////////

const D: i64 = 998244353;


//////////////////// 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()
}

// ax = by + 1 (a, b > 0)
fn linear_diophantine(a: i64, b: i64) -> Option<(i64, i64)> {
    if a == 1 {
        return Some((1, 0))
    }
    
    let q = b / a;
    let r = b % a;
    if r == 0 {
        return None
    }
    let (x1, y1) = linear_diophantine(r, a)?;
    Some((-q * x1 - y1, -x1))
}

fn inverse(a: i64, d: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, d).unwrap();
    if x >= 0 {
        x % d
    }
    else {
        x % d + d
    }
}


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

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

fn initialize_dp(X: usize) -> Vec<i64> {
    let mut dp: Vec<i64> = vec![0; X+1];
    dp[0] = 1;
    dp
}

fn update_dp(dp: &mut Vec<i64>, t: usize, tunes: &Vec<usize>) {
    let X = dp.len() - 1;
    let N = tunes.len() as i64;
    let rec_N = inverse(N, D);
    for &tune in tunes.iter() {
        let t1 = t + tune;
        if t1 <= X {
            dp[t1] = (dp[t1] + dp[t] * rec_N) % D
        }
    }
}

fn F(X: usize, tunes: Vec<usize>) -> i64 {
    let N = tunes.len() as i64;
    let rec_N = inverse(N, D);
    let mut dp = initialize_dp(X);
    let mut p: i64 = 0;
    for t in 0..X+1 {
        if t + tunes[0] > X {
            p = (p + dp[t] * rec_N) % D
        }
        update_dp(&mut dp, t, &tunes)
    }
    p
}

fn main() {
    let (X, tunes) = read_input();
    println!("{}", F(X, tunes))
}