546. Remove Boxes

先做出來一個backtracking的TLE版本，下面想想如何加memory， memory其實也是DP，狀態(tài)和轉(zhuǎn)移方程如下：
memo[l][r][3] 表示這樣狀態(tài)的解: [b_l, ..., b_r, A,A,A] with b_r==A
memo[l][r][k] = max(memo[l][r][k], memo[l][i][k+1] + memo[i+1][r-1][0])

Basically, if there is one i such that b_i==b_r, we partition the array into two: [b_l, ..., b_i, b_r, A, ..., A], and [b_{i+1}, ..., b_{r-1}]. The solution for first one will be memo[l][i][k+1], and the second will be memo[i+1][r-1][0]. Otherwise, we just remove the last k+1 boxes (including b_r) and search the best solution for lth to r-1th boxes

class Solution(object):
    def removeBoxes(self, boxes):
        """
        :type boxes: List[int]
        :rtype: int
        """
        if not boxes:
            return 0
        n = len(boxes)
        dp = [[[0]*n for _ in range(n)] for _ in range(n)]
        return self.search(dp, boxes, 0, n-1, 0)
    
    def search(self, dp, boxes, i, j, k):
        if i > j:
            return 0
        if dp[i][j][k] > 0:
            return dp[i][j][k]
        while j > i and boxes[j]==boxes[j-1]:
            j -= 1
            k += 1
        dp[i][j][k] = self.search(dp, boxes, i, j-1, 0) + (k+1)**2
        for m in range(i, j):
            if boxes[j] == boxes[m]:
                dp[i][j][k] = max(dp[i][j][k], self.search(dp, boxes, i, m, k+1) + self.search(dp, boxes, m+1, j-1, 0))
        return dp[i][j][k]