Leetcode 1937. Maximum Number of Points with Cost

Problem statement

https://leetcode.com/problems/maximum-number-of-points-with-cost/

Solution

The idea is to us dynamic programming, but we need to update each next row, using previous in linear time, we can not allow slower given problem constraints. Imagine, that we have row

A0, A1, A2, A3, A4, A5.

And we want to calcualte answers for the new row. For the first element we need to calculate: max(A0, A1-1, A2-2, A3-3, A4-4, A5-5), for the second element max(A0-1, A1, A2-1, A3-2, A4-3, A5-4), max(A0-2, A1-1, A2, A3-1, A4-2, A5-3) and so on.

The trick is to find all this values using some preprocessing. Look at the value max(A0-2, A1-1, A2, A3-1, A4-2, A5-3), it can be written as max(max([A0+0, A1+1, A2+2] - 2), max([A2-2, A3-3, A4-4, A5-5]) + 2). And in fact all values can be written in similar way. What we need to do now is to calculate cumulative maximums ans we are done!

Complexity

Time complexity is O(mn), space complexity is O(n).

Code

class Solution:
    def maxPoints(self, P):
        m, n = len(P), len(P[0])
        dp = P[0]
        for i in range(1, m):
            c1 = list(accumulate([a+b for a,b in zip(dp, range(n))], max))
            c2 = list(accumulate([a-b for a,b in zip(dp[::-1], range(n-1,-1,-1))], max))
            dp2 = [max(c1[i] - i, c2[n-1-i] + i) for i in range(n)]
            dp = [x+y for x,y in zip(dp2, P[i])]
        return max(dp)