Python, Straightforward [but slow] with Explanation

  • 3

    Let A be the array of boxes. A natural approach is to consider dp(i, j) to be the value of A[i: j+1]. Let's use a top-down dp and investigate this line of thinking.

    Say we want to evaluate dp(i, j). When we want to remove box A[i], we will do so with any subset of boxes of the same color within A[i: j+1]. Then our candidate answer will be (num_boxes_chosen)**2 + sum( dp(a, b) for [a,b] maximum size ranges not touching chosen boxes in [i,j] ). Since [a, b] has smaller size than [i, j], our dp function won't loop.

    Because in general a^2 + b^2 < (a+b)^2, we can consider contiguous ranges of "good" boxes (boxes that have the same color as A[i]) as one entity that we either include or don't include. Now we are ready for our algorithm:

    First, find the ranges of every good box. For example if A[1: 13] = [2, 2, 2, 3, 5, 6, 2, 2, 3, 2, 4, 5], our good ranges will be good = [[1, 3], [7, 8], [10]]. Then, for every subset of these good ranges, let's suppose we removed all the boxes represented by these subsets in one move. If for example we remove [1,3] + [10], then we've removed 4 boxes, and our candidate answer will be 4**2 + dp(1, 0) + dp(4,9) + dp(11,12). Our actual answer must be one of the candidate answers, so we'll take a running maximum.

    This approach is worst case exponential, so it's not the best approach but it passes the test cases.

    def removeBoxes(self, A):
        def outside_ranges(ranges, i, j):
            prev = i
            for r1, r2 in ranges:
                yield prev, r1 - 1
                prev = r2 + 1
            yield prev, j
        memo = {}
        def dp(i, j):
            if i >= j: return +(i==j)
            if (i,j) not in memo:
                good = []
                for k, v in itertools.groupby(range(i, j+1),
                        key = lambda x: A[x] == A[i]):
                    if k:
                        w = list(v)
                        good.append((w[0], w[-1]))
                ans = 0
                for size in xrange(1, len(good) + 1):
                    for subset in itertools.combinations(good, size):
                        cand = sum( g[-1] - g[0] + 1 for g in subset ) ** 2
                        cand += sum( dp(L, R) for L, R in outside_ranges(subset, i, j) )
                        ans = max(ans, cand)
                memo[i, j] = ans
            return memo[i, j]
        return dp(0, len(A)-1)

  • 0

    I'm a little confused- what constitutes a "good" box?

  • 1

    @jonesbp In our evaluation of dp(i, j), a good box is any box with the same color as A[i]. For example, if evaluating dp(1, 12), and A[1: 13] = [2, 2, 2, 3, 5, 6, 2, 2, 3, 2, 4, 5], then A[1], A[2], A[3], A[7], A[8], A[10] are the good boxes.

  • 2

    @awice Nice solution and implementation. I thought about a similar one, but decided that it wouldn't be accepted time-wise. Honestly, I am surprised it worked because it is exponential. It can't handle the input like [1,2,1,2,1,2,...]. You are lucky it is not in the test cases :)

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.