use std::cmp;

impl Solution {
    pub fn max_sum_after_partitioning(arr: Vec<i32>, k: i32) -> i32 {
        let mut dp: Vec<i32> = vec![0; arr.len() + 1];

        for i in 1..=arr.len() {
            let mut m = 0;
            let mut m_sum = 0;

            for j in (1..=k as usize).take_while(|j| i >= *j) {
                m = cmp::max(m, arr[i - j]);
                m_sum = cmp::max(m_sum, dp[i - j] + (j as i32) * m);
            }

            dp[i] = m_sum;
        }

        dp[arr.len()]
    }
}