LeetCode 546: Remove Boxes

Problem Restatement

We are given an array boxes.

Each integer represents the color of one box.

In one move, we can remove a continuous group of boxes with the same color. If the group has k boxes, we gain:

k * k

points.

After removing boxes, the remaining boxes close the gap. We repeat this until no boxes remain.

We need to return the maximum score possible.

The official problem states that each box color is represented by a positive number, and removing k continuous boxes of the same color gives k * k points. The constraint is 1 <= boxes.length <= 100.

Input and Output

Function shape:

def removeBoxes(boxes: list[int]) -> int:
    ...

Examples

Consider:

boxes = [1, 3, 2, 2, 2, 3, 4, 3, 1]

One strong strategy is:

remove [2, 2, 2] -> 3 * 3 = 9
remove [3, 3, 3] -> 3 * 3 = 9
remove [1, 1]    -> 2 * 2 = 4
remove [4]       -> 1 * 1 = 1

Total:

9 + 9 + 4 + 1 = 23

So the answer is:

23

Another example:

boxes = [1, 1, 1]

Remove all three together:

3 * 3 = 9

So the answer is:

9

For:

boxes = [1]

The answer is:

1

First Thought: Try Every Removal Order

A direct recursive idea is:

1. Find every continuous same-color segment.
2. Remove one segment.
3. Recursively solve the smaller array.
4. Take the best score.

This is correct, but too slow.

The same remaining box configuration can be reached through many different removal orders. Also, removing boxes changes adjacency, so the state is not just a simple index range.

We need dynamic programming.

Key Insight

The hard part is that boxes of the same color can become adjacent later.

Example:

[1, 2, 1]

If we remove the middle 2 first, the two 1 boxes become adjacent.

Then we can remove them together for:

2 * 2 = 4

This is better than removing them separately:

1 * 1 + 1 * 1 = 2

So a subproblem must remember not only the current interval, but also how many same-colored boxes are already waiting to merge with one side of the interval.

A standard state is:

dp(left, right, extra)

where dp(left, right, extra) means:

The maximum score obtainable from boxes[left:right + 1], assuming there are already extra boxes to the right of right with the same color as boxes[right].

Those extra boxes are not inside the interval, but they will be removed together with boxes[right].

Algorithm

Define:

dfs(left, right, extra)

If left > right, return 0.

Before solving the state, merge equal-colored boxes at the right end:

while left < right and boxes[right] == boxes[right - 1]:
    right -= 1
    extra += 1

This compression means that if the interval ends with several boxes of the same color, we treat them as one group with a larger extra.

Then we have two choices.

First, remove boxes[right] together with the extra boxes:

dfs(left, right - 1, 0) + (extra + 1) * (extra + 1)

Second, try to merge boxes[right] with an earlier box of the same color.

For each mid in [left, right - 1] where:

boxes[mid] == boxes[right]

we remove the middle interval first:

dfs(mid + 1, right - 1, 0)

Then boxes[mid] can merge with boxes[right] and the extra boxes:

dfs(left, mid, extra + 1)

Take the maximum over all choices.

Correctness

The state dfs(left, right, extra) represents all ways to remove the interval boxes[left:right + 1], with extra same-colored boxes attached to boxes[right] for a future removal.

For any optimal strategy for this state, consider what happens to boxes[right].

One possibility is that boxes[right] is removed without merging with any earlier box inside [left, right - 1]. Then it is removed together only with the extra boxes of the same color. The remaining interval is [left, right - 1], giving:

dfs(left, right - 1, 0) + (extra + 1) * (extra + 1)

The algorithm considers this case.

The other possibility is that boxes[right] is eventually removed together with some earlier box boxes[mid] of the same color. For that to happen, all boxes between mid and right must be removed first. That gives the subproblem:

dfs(mid + 1, right - 1, 0)

After that middle interval disappears, boxes[mid] becomes adjacent to boxes[right] and the extra boxes. So the remaining left interval becomes:

dfs(left, mid, extra + 1)

The algorithm considers every possible matching mid.

These cases cover every possible optimal strategy, because boxes[right] is either removed without merging with an earlier equal box, or it merges with at least one earlier equal box. The first such merge can be represented by some mid.

The recurrence takes the maximum score among all valid choices, so it computes the optimal value for each state.

The base case left > right is correct because an empty interval gives 0 points.

Therefore, dfs(0, n - 1, 0) returns the maximum score for the whole array.

Complexity

Let n = len(boxes).

The right-end compression reduces many repeated states in practice.

With n <= 100, this memoized interval DP is the intended approach.

Implementation

from functools import cache

class Solution:
    def removeBoxes(self, boxes: list[int]) -> int:
        n = len(boxes)

        @cache
        def dfs(left: int, right: int, extra: int) -> int:
            if left > right:
                return 0

            while left < right and boxes[right] == boxes[right - 1]:
                right -= 1
                extra += 1

            best = dfs(left, right - 1, 0) + (extra + 1) * (extra + 1)

            for mid in range(left, right):
                if boxes[mid] == boxes[right]:
                    best = max(
                        best,
                        dfs(mid + 1, right - 1, 0) + dfs(left, mid, extra + 1),
                    )

            return best

        return dfs(0, n - 1, 0)

Code Explanation

We use memoization:

@cache

because many recursive calls reach the same (left, right, extra) state.

The base case is:

if left > right:
    return 0

An empty interval has no score.

Then we compress same-colored boxes at the right edge:

while left < right and boxes[right] == boxes[right - 1]:
    right -= 1
    extra += 1

For example:

[1, 2, 3, 3, 3]

When the right edge is color 3, the three ending 3 boxes can be treated as one box plus extra = 2.

The direct removal option is:

best = dfs(left, right - 1, 0) + (extra + 1) * (extra + 1)

This removes the rightmost color group now.

Then we try delayed removal by merging with an earlier box of the same color:

for mid in range(left, right):
    if boxes[mid] == boxes[right]:

If boxes[mid] matches, we remove everything between mid and right first:

dfs(mid + 1, right - 1, 0)

Then boxes[mid] can merge with boxes[right], represented by:

dfs(left, mid, extra + 1)

The maximum over all choices is returned.

Small Walkthrough

For:

boxes = [1, 2, 1]

The direct strategy might remove the last 1 alone:

dfs(0, 1, 0) + 1

But the merge strategy notices that:

boxes[0] == boxes[2]

So it can remove the middle interval first:

dfs(1, 1, 0)

which scores 1.

Then it solves:

dfs(0, 0, 1)

meaning the first 1 has one extra 1 waiting to merge with it.

That removal scores:

(1 + 1) * (1 + 1) = 4

Total:

1 + 4 = 5

This is better than removing all boxes separately.

Testing

def run_tests():
    s = Solution()

    assert s.removeBoxes([1, 3, 2, 2, 2, 3, 4, 3, 1]) == 23
    assert s.removeBoxes([1, 1, 1]) == 9
    assert s.removeBoxes([1]) == 1
    assert s.removeBoxes([1, 2, 1]) == 5
    assert s.removeBoxes([1, 2, 2, 1]) == 8

    print("all tests passed")

run_tests()