rustgym/leetcode/
_546_remove_boxes.rs

1struct Solution;
2
3use std::collections::HashMap;
4
5impl Solution {
6    fn remove_boxes(boxes: Vec<i32>) -> i32 {
7        let n = boxes.len();
8        let mut memo: HashMap<(usize, usize, usize), usize> = HashMap::new();
9        Self::dp(0, n, 0, &mut memo, &boxes) as i32
10    }
11
12    fn dp(
13        mut start: usize,
14        end: usize,
15        mut k: usize,
16        memo: &mut HashMap<(usize, usize, usize), usize>,
17        boxes: &[i32],
18    ) -> usize {
19        if start == end {
20            return 0;
21        }
22        if let Some(&res) = memo.get(&(start, end, k)) {
23            return res;
24        }
25        while start + 1 < end && boxes[start + 1] == boxes[start] {
26            start += 1;
27            k += 1;
28        }
29        let mut res = Self::dp(start + 1, end, 0, memo, boxes) + (k + 1) * (k + 1);
30        for i in start + 1..end {
31            if boxes[i] == boxes[start] {
32                res = res.max(
33                    Self::dp(i, end, k + 1, memo, boxes) + Self::dp(start + 1, i, 0, memo, boxes),
34                );
35            }
36        }
37        memo.insert((start, end, k), res);
38        res
39    }
40}
41
42#[test]
43fn test() {
44    let boxes = vec![1, 3, 2, 2, 2, 3, 4, 3, 1];
45    let res = 23;
46    assert_eq!(Solution::remove_boxes(boxes), res);
47}
rustgym/leetcode/_546_remove_boxes.rs

rustgym/leetcode/
_546_remove_boxes.rs
