https://leetcode.com/problems/best-time-to-buy-and-sell-stock-iii

This problem is a special case of problem 188. Best Time to Buy and Sell Stock IV. The following solution works for both problems. The only difference is that here we have k = 2 transactions and in problem 188 we can have different k.

My way to solve this problem is to first evaluate matrix B of differences and then what is asked is to find the maximum sum of two(k) contiguous subarrays. We can also make optimization if k > len(prices)//2: return sum(x for x in B if x > 0), which will help us, if k is big (in this problem it is equal to 2, so you can remove this line, but you need it for problem 188). If k is more than half of length of prices, we can just choose all positive elements, we will have enough trancastions to do it.

Let us create dp array with size k+1 by n-1, where dp[i][j] is the maximum gain, where already made j transactions, that is choose j contiguous subarrays and used all elements before our equal number i. Also, mp[i][j] = max(dp[0][j], ..., dp[i][j]). We take k+1 size, because we start with 0 transactions, which will be filled with zeros. We take n-1 size, because original size is n, and size of differences is n-1. Also we start with dp[0][1] = B[0], because we need to choose one contiguous subarray which ends with element B[0], which is B[0] itself. Also we put mp[0][1] = B[0] for the same logic.

Now, about updates: we iterate over all i from 1 to n-2 inclusive and j from 1 to k inclusive and:

  1. Update dp[i][j] = max(mp[i-1][j-1], dp[i-1][j]) + B[i]. By definition we need to take B[i]. We can either say, that we add it to the last contiguous subarray: dp[i-1][j] + B[i], or we say that it is new contiguous subarray: mp[i-1][j-1] + B[i]. Note, here we use mp, because we actually have max(dp[0][j-1], ... , dp[i-1][j-1]).
  2. Update mp[i][j] = max(dp[i][j], mp[i-1][j]).
  3. Finally, return maximum from the mp[-1], we need to choose maximum, because optimal solution can be with less than k transactions.

Complexity: Time complexity is O(nk) = O(n), because here k = 2. Space complexity is also O(nk) = O(n).

class Solution:
    def maxProfit(self, prices):
        if len(prices) <= 1: return 0
        n, k = len(prices), 2

        B = [prices[i+1] - prices[i] for i in range(len(prices) - 1)]
        if k > len(prices)//2: return sum(x for x in B if x > 0)
        
        dp = [[0]*(k+1) for _ in range(n-1)] 
        mp = [[0]*(k+1) for _ in range(n-1)] 

        dp[0][1], mp[0][1] = B[0], B[0]

        for i in range(1, n-1):
            for j in range(1, k+1):
                dp[i][j] = max(mp[i-1][j-1], dp[i-1][j]) + B[i]
                mp[i][j] = max(dp[i][j], mp[i-1][j])

        return max(mp[-1])

Optimization:

Note, that if we have array like [1,5,7, -7, -4, -3, 10, 2, 7, -4, -8, 13, 15], then we can work in fact with smaller array [1+5+7, -7-4-3, 10+2+7, -4-8, 13+15] = [13,-14,19,-12,28]. So, instead of B = [prices[i+1] - prices[i] for i in range(len(prices) - 1)], we can evaluate:

delta = [prices[i+1]-prices[i] for i in range (len(prices)-1)]
B=[sum(delta) for _, delta in groupby(delta, key=lambda x: x < 0)]
n, k = len(B) + 1, 2

When used this code, I have 132ms to 64ms improvement (which is faster than 99%). Thanks user2349 for providing shorter version of code!

If you like the solution, you can upvote it on leetcode discussion section: Problem 0123