546. Remove Boxes

• Time: $O(n^4)$
• Space: $O(n^3)$

C++JavaPython

class Solution {
 public:
  int removeBoxes(vector<int>& boxes) {
    const int n = boxes.size();
    vector<vector<vector<int>>> mem(n, vector<vector<int>>(n, vector<int>(n)));
    return removeBoxes(boxes, 0, n - 1, 0, mem);
  }

 private:
  // Returns the maximum score of boxes[i..j] if k boxes eqaul to boxes[j].
  int removeBoxes(const vector<int>& boxes, int i, int j, int k,
                  vector<vector<vector<int>>>& mem) {
    if (i > j)
      return 0;
    if (mem[i][j][k] > 0)
      return mem[i][j][k];

    int r = j;
    int sameBoxes = k + 1;
    while (r > 0 && boxes[r - 1] == boxes[r]) {
      --r;
      ++sameBoxes;
    }
    mem[i][j][k] = removeBoxes(boxes, i, r - 1, 0, mem) + sameBoxes * sameBoxes;

    for (int p = i; p < r; ++p)
      if (boxes[p] == boxes[r])
        mem[i][j][k] =
            max(mem[i][j][k], removeBoxes(boxes, i, p, sameBoxes, mem) +
                                  removeBoxes(boxes, p + 1, r - 1, 0, mem));

    return mem[i][j][k];
  }
};

class Solution {
  public int removeBoxes(int[] boxes) {
    final int n = boxes.length;
    int[][][] mem = new int[n][n][n];
    return removeBoxes(boxes, 0, n - 1, 0, mem);
  }

  // Returns the maximum score of boxes[i..j] if k boxes are equal to boxes[j]
  private int removeBoxes(int[] boxes, int i, int j, int k, int[][][] mem) {
    if (i > j)
      return 0;
    if (mem[i][j][k] > 0)
      return mem[i][j][k];

    int r = j;
    int sameBoxes = k + 1;
    while (r > 0 && boxes[r - 1] == boxes[r]) {
      --r;
      ++sameBoxes;
    }
    mem[i][j][k] = removeBoxes(boxes, i, r - 1, 0, mem) + sameBoxes * sameBoxes;

    for (int p = i; p < r; ++p)
      if (boxes[p] == boxes[r])
        mem[i][j][k] = Math.max(mem[i][j][k], removeBoxes(boxes, i, p, sameBoxes, mem) +
                                                  removeBoxes(boxes, p + 1, r - 1, 0, mem));

    return mem[i][j][k];
  }
}

class Solution:
  def removeBoxes(self, boxes: list[int]) -> int:
    @functools.lru_cache(None)
    def dp(i: int, j: int, k: int) -> int:
      """
      Returns the maximum score of boxes[i..j] if k boxes equal to boxes[j].
      """
      if i > j:
        return 0

      r = j
      sameBoxes = k + 1
      while r > 0 and boxes[r - 1] == boxes[r]:
        r -= 1
        sameBoxes += 1
      res = dp(i, r - 1, 0) + sameBoxes * sameBoxes

      for p in range(i, r):
        if boxes[p] == boxes[r]:
          res = max(res, dp(i, p, sameBoxes) + dp(p + 1, r - 1, 0))

      return res

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