Maximum Profit from K Transactions

From LeetCode:

You have an array for which the i-th element is the price of a given stock on day i. Design an algorithm to find the maximum profit. You may complete at most k transactions.

Note: You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).

Input: prices = [3, 2, 6, 5, 0, 3], k = 2
Output: 7
# Explanation:
# 1. Buy on day 2 and sell on day 3, profit = 6 - 2 = 4.
# 2. Buy on day 5 and sell on day 6, profit = 3 - 0 = 3.

The first thing to note about this question is that allows for up to k transactions –– which means we can transact fewer times if necessary in order to maximise profit. This implies that we should assess profits for 1, 2, 3… and all the way up to k transactions. Once again, our instinct for solving subproblems naturally leads us to consider Dynamic Programming (DP) as a possible course of action.

Dynamic Programming in Broad Strokes

We’ve covered a few DP problems so far, but we haven’t talked about how we can develop a mental model that applies more broadly to different types of questions.

Many who’ve studied computer science formally will be familiar with the two defining characteristics of DP: Optimal Substructure and Overlapping Subproblems. The central idea is simple: We can find the answer to a complex problem by 1) starting from the smallest, simplest version of it and 2) gradually increasing the problem space.

Before using DP to tackle a problem, get some clarity on the subproblems you’ll be addressing. In particular, consider the dimensions of the problem and see whether you can decrement them. To use a couple of examples:

• In the knapsack problem, we want to retrieve the most amount of value from a room full of items but are constrained by the amount of weight we can carry. What are the dimensions involved? The first would be the set of items available to us. The second would be the weight we’re able to tolerate. Our smallest subproblem? An empty room and a weight tolerance of zero. We work our way up to the complete solution by incrementally adjusting both parameters.
• Earlier, we tackled the Partition Array into K Equal Sum Subsets problem. There’s really just one key dimension involved –– the different types of subset configurations. In our algorithm, we sought to understand larger subsets by using the information that we had gathered on smaller ones.

Let’s now apply these same principles to our profit-maximisation problem.

Tracking Best Attainable Profits

We’ll use a profits array to keep track of our best profits. To examine all profit possibilities, we’ll have to incrementally increase the constraint on the number of transactions.

def max_profit(k, prices):
    if not k or not prices:
        return 0

    n = len(prices)

    # `profits[i]` tracks the best profits attainable when
    # we *add* the option of selling at `i`
    profits = [0 for _ in range(n)]

    # `txn_idx` describes the currents transaction count
    for txn_idx in range(1, k+1):

        # `new_profits` tracks the best profits attainable when
        # we *increment* the number of allowable transactions
        new_profits = profits[:]
        # ...

        for sell_idx in range(1, n):
            # ...  

Note that the value in profits[i] isn’t always the result of a sale on day i. Rather, the value in profits[i] can only be said to account for the possibility of selling on day i. Consider the following:

prices = [2, 8, 5]
# txn_idx = 1, profits = [0, 6, 6]

At txn_idx = 1, when we compute the value of profits[2] , we account for the possibility of selling at prices[2] = 5 but don’t actually act on it because this yields us less profit compared to selling at prices[1] = 8. Subsequently, we carry over the computed best profits from profits[1] to profits[2].

Tracking the Best Preceding State Before Selling

Besides the need to evaluate every selling price, there’s also a need to review:

• The profit we’ve made from previous transactions
• The buying price we’ve chosen

To maximise profits, we’ll want to optimise for these two factors such that their combined sum is as high as possible. Let’s represent their sum with a single term: best_preceding_value.

profits = [0 for _ in range(n)]

for txn_idx in range(1, k+1):
    
    new_profits = profits[:]

    # best_preceding_value =
    # best profits before purchase - best_buy_price
    best_preceding_value = -prices[0]

    for sell_idx in range(1, n):
        # ...

best_preceding_value is a term that’s critical to how we judge every possible sale. It tells us: Before making a sale at sell_idx, what’s our best possible financial state? We’ve used two separate arrays for profit tracking (profits vs new_profits) because best_preceding_value is determined by the profits from preceding transactions. Specifically, we’ll need to know what we gained previously from making up to txn_idx-1 transactions.

Evaluating Selling Prices

At each sell_idx, the question we should ask is this: Does executing a sale at this selling price top the maximum profits we’ve earned from transacting up to sell_idx-1? If yes, we’ll update prices[sell_idx] to store this improved profit outcome. Otherwise, we’ll just carry over the profits we made in prices[sell_idx-1].

profits = [0 for _ in range(n)]

for txn_idx in range(1, k+1):
    
    new_profits = profits[:]
    best_preceding_value = -prices[0]

    for sell_idx in range(1, n):
        # Evaluate selling price
        new_profits[sell_idx] = max(
            new_profits[sell_idx-1],
            best_preceding_value + prices[sell_idx]
        )

Once we’ve updated prices[sell_idx], we’ll need to also update best_preceding_value so that we can properly evaluate the selling price at the next index (i.e. sell_idx+1).

To do this, we’ll consider 1) prices[sell_idx] as a buying price and 2) the best preceding profits we earned from txn_idx-1 transactions, up to sell_idx-1:

profits = [0 for _ in range(n)]

for txn_idx in range(1, k+1):
    
    new_profits = profits[:]
    best_preceding_value = -prices[0]

    for sell_idx in range(1, n):
        new_profits[sell_idx] = max(
            new_profits[sell_idx-1],
            best_preceding_value + prices[sell_idx]
        )

        # Update best_preceding_value
        best_preceding_value = max(
            best_preceding_value,
            profits[sell_idx-1] - prices[sell_idx]
        )

    profits = new_profits

Once we’ve iterated through all selling prices across k transactions, we’ll find the value of our best attainable profit sitting nicely in profits[-1].

Given n prices, our algorithm runs in O(n) space and O(nk) time.

With that, we’ve solved the Maximum Profit from K Transactions problem 📈💰.

Full Solution

def max_profit(k, prices):
    if not k or not prices:
        return 0

    n = len(prices)

    # (Optional) Additional optimisation if we
    # are able to execute unlimited transactions
    if k >= len(prices) * 2:
        max_profits = 0
        for i in range(1, n):
            max_profits += max(prices[i] - prices[i-1], 0)
        return max_profits


    # Core DP algorithm
    profits = [0 for _ in range(n)]

    for txn_idx in range(1, k + 1):
        new_profits = profits[:]
        best_preceding_value = -prices[0]

        for sell_idx in range(1, n):
            new_profits[sell_idx] = max(
                new_profits[sell_idx-1],
                best_preceding_value + prices[sell_idx]
            )

            best_preceding_value = max(
                best_preceding_value,
                profits[sell_idx-1] - prices[sell_idx]
            )

        profits = new_profits

    return profits[-1]