アルゴリズムと数学 020

https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_t

選んだ枚数とカードの数の和を状態としたDPですね。

// Choose Cards 1
#![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<i32> {
    let _N: usize = read();
    let A = read_vec();
    A
}

fn update(dp0: Vec<Vec<i32>>, a: i32) -> Vec<Vec<i32>> {
    let mut dp = (0..6).map(|i| dp0[i].to_vec()).collect::<Vec<Vec<i32>>>();
    for i in 0usize..5 {
        for j in 0..((1001-a) as usize) {
            dp[i+1][j+(a as usize)] += dp0[i][j]
        }
    }
    dp
}

fn f(A: Vec<i32>) -> i32 {
    let mut dp = (0..6).map(|_| (0..1001).map(|_| 0).collect::<Vec<i32>>()).
                                                    collect::<Vec<Vec<i32>>>();
    dp[0][0] = 1;
    for a in A.into_iter() {
        dp = update(dp, a)
    }
    dp[5][1000]
}

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