546. Remove Boxes

You are given several boxes with different colors represented by different positive numbers.

You may experience several rounds to remove boxes until there is no box left. Each time you can choose some continuous boxes with the same color (i.e., composed of k boxes, k >= 1), remove them and get k * k points.

Return the maximum points you can get.

 

Example 1:

Input: boxes = [1,3,2,2,2,3,4,3,1]
Output: 23
Explanation:
[1, 3, 2, 2, 2, 3, 4, 3, 1] 
----> [1, 3, 3, 4, 3, 1] (3*3=9 points) 
----> [1, 3, 3, 3, 1] (1*1=1 points) 
----> [1, 1] (3*3=9 points) 
----> [] (2*2=4 points)

Example 2:

Input: boxes = [1,1,1]
Output: 9

Example 3:

Input: boxes = [1]
Output: 1

 

Constraints:

  • 1 <= boxes.length <= 100
  • 1 <= boxes[i] <= 100

Rust Solution

struct Solution;

use std::collections::HashMap;

impl Solution {
    fn remove_boxes(boxes: Vec<i32>) -> i32 {
        let n = boxes.len();
        let mut memo: HashMap<(usize, usize, usize), usize> = HashMap::new();
        Self::dp(0, n, 0, &mut memo, &boxes) as i32
    }

    fn dp(
        mut start: usize,
        end: usize,
        mut k: usize,
        memo: &mut HashMap<(usize, usize, usize), usize>,
        boxes: &[i32],
    ) -> usize {
        if start == end {
            return 0;
        }
        if let Some(&res) = memo.get(&(start, end, k)) {
            return res;
        }
        while start + 1 < end && boxes[start + 1] == boxes[start] {
            start += 1;
            k += 1;
        }
        let mut res = Self::dp(start + 1, end, 0, memo, boxes) + (k + 1) * (k + 1);
        for i in start + 1..end {
            if boxes[i] == boxes[start] {
                res = res.max(
                    Self::dp(i, end, k + 1, memo, boxes) + Self::dp(start + 1, i, 0, memo, boxes),
                );
            }
        }
        memo.insert((start, end, k), res);
        res
    }
}

#[test]
fn test() {
    let boxes = vec![1, 3, 2, 2, 2, 3, 4, 3, 1];
    let res = 23;
    assert_eq!(Solution::remove_boxes(boxes), res);
}

Having problems with this solution? Click here to submit an issue on github.