Python O(nk)/O(n) clean solution

  • 2

    This DP solution is well explained on:
    Although I used O(n) space (we only need to keep track of two n-length arrays in each iteration) with O(nk) time complexity, I still got TLE in one of the tests where k >= len(prices)/2. In this case, the problem is conveniently reduced to stock II. So I reused the function I wrote for that to handle the egde case:

    class Solution(object):
        def maxProfit(self, k, prices):
            if not k or not prices: return 0
            l = len(prices)
            if k >= l/2: return self.stock2(prices)
            f  = [[0 for _ in range(l)],[0 for _ in range(l)]]
            for i in range(1, k+1):
                tmpmax = - prices[0]
                for j in range(1, l):
                    f[1][j] = max(f[1][j-1], prices[j] + tmpmax)
                    tmpmax = max(tmpmax, f[0][j] - prices[j])
                f[0], f[1] = f[1], [0 for _ in range(l)]
            return f[0][-1]
        def stock2(self, prices):
            maxprof = 0
            for i, p in enumerate(prices):
                if i != 0 and prices[i] - prices[i-1] > 0:
                    maxprof += prices[i] - prices[i-1]
            return maxprof

  • 0

    Nice solution. Just refine the code a little bit.

    def maxProfit(k, prices):
        n = len(prices)
        if k > n/2: # if k is too large, all potential profit can be earned
            return sum([prices[i] - prices[i-1] for i in range(1, n) \
                if prices[i] > prices[i-1] ])
        f, fnew = [0] * n, [0] * n  # ping-pong dp, maximum profit can get then
        for _ in range(k):
            tmax = - prices[0]  # the money I have today
            for i in range(1, n):
                fnew[i] = max(fnew[i-1], prices[i] + tmax) # do nothing or sell one
                tmax = max(tmax, f[i] - prices[i]) # max money after this day
            f, fnew = fnew, [0] * n
        return f[n-1]

Log in to reply

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