Remove Boxes

Intuition

Solution

class Solution:
    def removeBoxes(self, boxes: List[int]) -> int:
        self.cache = {}
        self.boxes = boxes
        return self.dp(0, len(boxes), 0)

    def dp(self, i, j, k):
        if (i, j, k) in self.cache:
            return self.cache[(i, j, k)]

        while i < j - 1 and self.boxes[i] == self.boxes[i + 1]:
            i += 1
            k += 1

        if i == (j - 1):
            return (k + 1) * (k + 1)

        ans = self.dp(i + 1, j, 0) + (k + 1) * (k + 1)
        for m in range(i + 1, j):
            if self.boxes[m] != self.boxes[i]:
                continue
            ans = max(ans, self.dp(i + 1, m, 0) + self.dp(m, j, k + 1))
        self.cache[(i, j, k)] = ans
        return ans