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 [ ]: