LeetCode/rs/maximum-value-of-k-coins-from-piles.rs

use std::cmp;

struct Solution {}
impl Solution {
    pub fn max_value_of_coins(piles: Vec<Vec<i32>>, k: i32) -> i32 {
        let k = k as usize;

        let mut mem: Vec<Vec<i32>> = vec![];
        mem.resize_with(piles.len() + 1, || {
            let mut v = vec![];
            v.resize(k as usize + 1, 0);
            v
        });

        for i in 1..=piles.len() {
            for max_j in 0..=k {
                let mut running_sum = 0;

                for j in 0..=cmp::min(piles[i - 1].len(), max_j as usize) {
                    if j > 0 {
                        running_sum += piles[i - 1][j - 1];
                    }
                    mem[i][max_j] = cmp::max(mem[i][max_j], running_sum + mem[i - 1][max_j - j]);
                }
            }
        }

        mem[piles.len()][k]
    }
}

fn main() {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_1() {
        let piles = vec![vec![1, 100, 3], vec![7, 8, 9]];
        let expected = vec![0, 7, 101, 108, 116, 125, 128, 128];

        for (k, right) in expected.iter().enumerate() {
            assert_eq!(
                Solution::max_value_of_coins(piles.clone(), k as i32),
                *right
            );
        }
    }

    #[test]
    fn example_2() {
        let piles = vec![
            vec![100],
            vec![100],
            vec![100],
            vec![100],
            vec![100],
            vec![100],
            vec![1, 1, 1, 1, 1, 1, 700],
        ];
        let expected = vec![
            0, 100, 200, 300, 400, 500, 600, 706, 806, 906, 1006, 1106, 1206, 1306, 1306,
        ];

        for (k, right) in expected.iter().enumerate() {
            assert_eq!(
                Solution::max_value_of_coins(piles.clone(), k as i32),
                *right
            );
        }
    }
}