In [78]:
test1 = [9,8,7,6,5,4,3,2] # profit_k = 0 for all k
test2 = [9, 1, 10, 13, 0] # profit_1 = 12, profit_2 = 13
test3 = [0,5,1,10,2,20,-8] # profit_1 = 20, profit_2 = 29
test4 = [1,2,3,4,5] # profit_1 = 4, profit_2 = 6
In [47]:
def maxProfit_1(prices):
    # returns max profit if at most one transaction is allowed
    # linear time
    
    mins = {} # i -> min(prices[:i+1])
    mins[0] = prices[0] 
    
    def maxProfit(prices, i): 
        # returns max profit if at most one transaction is allowed, and sell at index >= i
    
        if i >= len(prices):
            return 0

        if prices[i] < mins[i-1]: # new min
            mins[i] = prices[i]
            return maxProfit(prices, i+1)
            
        mins[i] = mins[i-1]   
        sell = prices[i] - mins[i] # profit if sell i th entry
        not_sell = maxProfit(prices, i+1) # profit if not sell i th entry
        return max(sell, not_sell)
                
    return maxProfit(prices, 1)
In [48]:
maxProfit_1(test3)
Out[48]:
20
In [87]:
def maxProfit(prices, k):
    # returns max profit if at most k transactions are allowed in O(len(prices)*k^2) time

    memo = {}
    
    def value(prices, i, b, s) -> int: 
        # returns max additional profit that can be made where b <= s <= k is the remaining number of things that may be bought, sold resp.  
        # and we can still do transactions at indices >= i. 
    
        if i >= len(prices):
            return 0
            
        if s == 0: # we already did max number of transactions
            return 0

        if (i+1, b, s) not in memo:
            memo[(i+1, b, s)] = value(prices, i+1, b, s) 
            
        out = memo[(i+1, b, s)] # do nothing at index i
            
        if b > 0:
            if (i+1, b-1, s) not in memo:
                memo[(i+1, b-1, s)] = value(prices, i+1, b-1, s) 
            buy = -prices[i] + memo[(i+1, b-1, s)] #buy at index i

            if buy > out:
                out = buy

        if s > b:
            if (i+1, b, s-1) not in memo:
                memo[(i+1, b, s-1)] = value(prices, i+1, b, s-1) 
            sell = prices[i] + memo[(i+1, b, s-1)]

            if sell > out:
                out = sell

        return out
        
    return value(prices, 0, k, k)
    
In [88]:
maxProfit(test3, 2)
Out[88]:
29
In [ ]:
 
In [ ]: