BinarySearch 0632 Maximum Points From Removals

Problem statement

https://binarysearch.com/problems/Maximum-Points-From-Removals/

Solution

The same as Leetcode 546. Remove Boxes.

Complexity

It is O(n^4) for time and O(n^3) for space. I think more careful analysis can be made.

Code

class Solution:
    def solve(self, B):
        @lru_cache(None)
        def dp(i, j, k):
            if i > j: return 0
            indx = [m for m in range(i+1, j+1) if B[m] == B[i] and (m == i+1 or B[m] != B[m-1])]
            ans = (k+1)**2 + dp(i+1, j, 0)
            return max([ans] + [dp(i+1, m-1, 0) + dp(m, j, k+1) for m in indx])   

        return dp(0, len(B)-1, 0)