Skip to content

546. Remove Boxes 👍

  • Time: $O(n^4)$
  • Space: $O(n^3)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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];
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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];
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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)