leetcode-546. 移除盒子

题目

给出一些不同颜色的盒子，盒子的颜色由数字表示，即不同的数字表示不同的颜色。
你将经过若干轮操作去去掉盒子，直到所有的盒子都去掉为止。每一轮你可以移除具有相同颜色的连续 k 个盒子（k >= 1），这样一轮之后你将得到 k*k 个积分。
当你将所有盒子都去掉之后，求你能获得的最大积分和。

示例：

输入：boxes = [1,3,2,2,2,3,4,3,1]
输出：23
解释：
[1, 3, 2, 2, 2, 3, 4, 3, 1] 
----> [1, 3, 3, 4, 3, 1] (3*3=9 分) 
----> [1, 3, 3, 3, 1] (1*1=1 分) 
----> [1, 1] (3*3=9 分) 
----> [] (2*2=4 分)

提示：

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

解题思路

暴力版：
从前向后移除盒子，然后再对剩下的盒子递归，时间复杂度是$o(n!)$，超时了

代码

暴力版：

class Solution:
    def removeBoxes(self, boxes: List[int]) -> int:

        def merge_boxes(box_cnts: tuple) -> tuple:
            ans = []
            prev = None
            prev_cnt = 0
            for pair in box_cnts:
                if not prev:
                    prev = pair[0]
                    prev_cnt = pair[1]
                    continue
                if prev != pair[0]:
                    ans.append((prev, prev_cnt))
                    prev, prev_cnt = pair
                else:
                    prev_cnt += pair[1]
            ans.append((prev, prev_cnt))
            return tuple(ans)

        @lru_cache(None)
        def helper(box_cnts: tuple) -> int:
            if not box_cnts:
                return 0
            if len(box_cnts) == 1:
                return box_cnts[0][1] ** 2
            merged_tuple = merge_boxes(box_cnts)
            return max(item[1] ** 2 + helper(merged_tuple[:index] + merged_tuple[index + 1:]) for index, item in enumerate(merged_tuple))

        box_counts = []
        prev = None
        cnt = 0
        for num in boxes:
            if not prev:
                prev = num
                cnt += 1
                continue
            if num != prev:
                box_counts.append((prev, cnt))
                prev = num
                cnt = 0
            cnt += 1
        box_counts.append((prev, cnt))

        return helper(tuple(box_counts))