Stock Market Shark: More Multidimensional DP
Today will be the final problem we do (for now) comparing Rust and Haskell LeetCode solutions. We’ll do a wrap-up of some of the important lessons next week. Last week’s problem was a multi-dimensional dynamic programming problem where the “dimensions” were obvious. We were working in 2D space trying to find the largest square, so we wanted the cells in our “DP” grid to correspond to the cells in our input grid.
Today we’ll solve one final problem using DP in multiple dimensions where the dimensions aren’t quite as obvious. To learn more about the basics behind implementing DP in Haskell, you need to enroll in our course, Solve.hs! You’ll learn many principles about algorithms in Module 3 and get a ton of practice with our exercises!
The Problem
Today’s problem is Best Time to Buy and Sell Stack IV, the final in a series of problems where we are aiming to maximize the profit we can make from purchasing a single stock.
We have two problem inputs. The first is an array of the prices of the stock over a number of days. Each day has one price. There is no fluctuation over the course of a day (real world stock trading would be much easier if we got this kind of future data!).
Our second input is a number of “transactions” we can make. A single transaction consists of buying AND selling the stock. There are some restrictions on how these transactions work. The primary one is that we cannot have simultaneous transactions. Another way of saying this is that we can only hold one “instance” of the stock at a time. We can’t buy one instance of the stock on day 1, and then another instance on day 2, and then sell them both later.
We also cannot sell a stock on the same day we buy it, nor buy a new instance on the same day we sell a previous instance. This isn’t so much a problem constraint as an algorithmic insight that there is no benefit to us doing this. Buying and selling on the same day yields no net profit, so we may as well just not use the transaction.
As an example, suppose we have 3 transactions to use, and the following data for the upcoming days:
[1, 4, 8, 2, 7, 1, 15]The solution here is 26, via the following transactions:
- Buy the stock on day 1 for $1, sell it on day 3 for $8 ($7 profit)
- Buy the stock on day 4 for $2, sell it on day 5 for $7 ($5 profit)
- Buy the stock on day 6 for $1, sell it on day 7 for $15 ($14 profit)
If we only had 2 transactions to work with, the answer would be 21. We would simply omit the second transaction.
The Algorithm
Since this is a “hard” problem, the algorithm description is a bit tricky! But we can break it into a few pieces.
Grid Structure
As I alluded to, this is a multi-dimensional DP problem, but the “dimensions” are not as clear as our last problem, because this problem doesn’t have a spatial nature. But once you do enough DP problems, it gets easier to see what the dimensions are.
One dimension will be the “current day”, and the other will be the “transaction state”. The cell {s, d} will indicate “Given I am in state s on day d, what is the largest additional profit I can achieve?”
The number of days is obviously equal to the size of our input array. This will be our column dimension. So column i will always mean “if I am in this state on day i”.
The number of transaction states is actually double the number of transactions we are allowed. We want one row for each transaction to capture the state after we have bought for this transaction, and one row for before buying as part of this transaction (we’ll refer to this row as “pre-bought” throughout).
We’ll order the rows so that earlier rows represent fewer transactions remaining. Thus the first row indicates the state of having purchased the stock for the final transaction, but not yet having sold it. The second row indicates you have one transaction still available, but you haven’t bought the stock for this transaction yet. The third row indicates you have purchased the stock and you’ll have 1 complete transaction remaining after selling it. And so on. So with n days and k transactions, our grid will have size 2k x n.
Base Cases
Now let’s think about the base cases of this grid. It is easiest to consider the last day, the final column of the grid. If we’re on the last day, the marginal gain we can make if we are holding the stock is simply to sell it (all prices are positive), which would give us a “profit” of the final sale price. We don’t need to consider the cost of buying the stock for these rows. We just think about “given that I have the stock, what’s the most I can end up with”.
Then, for all the “pre-bought” rows, the final column is 0. We don’t have enough time to buy AND sell a stock, so we just do nothing.
Now we can also populate the rows for the final transaction fairly easily. These are base cases as well. We’ll populate them from right to left, meaning from the later days to the earlier days (recall we’ve already filled in the very last day).
For the “top” row, where we’ve already bought the stock for our final transaction, we have two choices. We can “sell” the stock on that day, or “keep” the stock to cell later. The first option means we just use the price for that day, and the second means we use the recorded value for the next day. We want the maximum of these options.
Once we’ve populated the “bought” row, we move on to the “pre-bought” row below it. Again, we’ll loop right to left and have two options each time. We can “buy” the stock, which would move us “up” to the bought row on the next day, except we have to subtract the price of the stock. Or we can “stay” and not buy the stock. This means we grab the value from the same row in the next column. Again, we just use the max of these two options.
At this point, we’ve populated the entire last column of our grid AND the first two rows.
Recursive Cases
For the “recursive” cases (we can actually think of them as “inductive” cases), we go two rows at a time, counting up to our total transaction count. Each transaction follows the same pattern, which is similar to what we did for the rows above.
First, fill in the “bought” row for this transaction. We can “sell” or “keep” the stock. Selling moves us up and to the right, and adds the sale price for that day. But keeping moves us directly right in our grid. Again, we take the max of these options.
Then we fill the “pre-bought” row for this transaction. We can “buy” or “stay”. Buying means subtracting the price for that day from the value up and to the right. Staying means we take the value immediately to our right. As always, take the max.
When we’ve completed populated our grid following this pattern, our final answer is the value in the bottom left of the grid! This is the maximum profit starting from day 0 and before buying for any of our transactions, which is the true starting state of the problem.
Rust Solution
Let’s solve this in Rust first! We begin by defining a few values and handling an edge case (if there’s only 1 day, the answer is 0 since we can’t buy and sell).
pub fn max_profit(k: i32, prices: Vec<i32>) -> i32 {
let n = prices.len();
if n == 1 {
return 0;
}
let ku = k as usize;
let numRows = 2 * ku;
// Create our zero-ed out grid
let mut dp: Vec<Vec<i32>> = Vec::with_capacity(numRows);
dp.resize(numRows, Vec::with_capacity(n));
for i in 0..numRows {
dp[i].resize(n, 0);
}
...
}Now we handle the first two rows (our “final” transaction). In each case, we start with the base case of the final day, and then move from right to left, following the rules described in the algorithm.
pub fn max_profit(k: i32, prices: Vec<i32>) -> i32 {
...
// Final Transaction
// Always sell on the last day!
dp[0][n - 1] = prices[n - 1];
for i in (0..=(n-2)).rev() {
// Sell or Keep
dp[0][i] = std::cmp::max(prices[i], dp[0][i+1]);
}
dp[1][n - 1] = 0;
for i in (0..=(n-2)).rev() {
// Buy (subtract price!) or keep
dp[1][i] = std::cmp::max(dp[0][i+1] - prices[i], dp[1][i+1]);
}
...
}Now we write our core loop, going through the remaining transaction count. We start by defining the correct row numbers and setting the final-column base cases:
pub fn max_profit(k: i32, prices: Vec<i32>) -> i32 {
// Setup
...
// Final Transaction
...
// All other transactions
for j in 1..ku {
let boughtRow = 2 * j;
let preBoughtRow = boughtRow + 1;
// Always sell on the last day!
dp[boughtRow][n - 1] = prices[n - 1];
// 0 - No time to buy/sell!
dp[preBoughtRow][n - 1] = 0;
...
}
}And now we apply the logic for our algorithm. As we populate each row from right to left, we simply apply our two choices: sell/keep for the “bought” row and buy/stay for the “pre-bought” row.
pub fn max_profit(k: i32, prices: Vec<i32>) -> i32 {
...
// All other transactions
for j in 1..ku {
let boughtRow = 2 * j;
let preBoughtRow = boughtRow + 1;
// Always sell on the last day!
dp[boughtRow][n - 1] = prices[n - 1];
// 0 - No time to buy/sell!
dp[preBoughtRow][n - 1] = 0;
// Sell or Keep!
for i in (0..=(n-2)).rev() {
dp[boughtRow][i] = std::cmp::max(dp[boughtRow - 1][i+1] + prices[i], dp[boughtRow][i + 1]);
}
// Buy or Stay!
for i in (0..=(n-2)).rev() {
dp[preBoughtRow][i] = std::cmp::max(dp[boughtRow][i+1] - prices[i], dp[preBoughtRow][i + 1])
}
}
return dp[numRows - 1][0];
}This completes our loop, and the final thing we need, as you can see, is to return the value in the bottom left of our grid!
Here is the complete solution:
pub fn max_profit(k: i32, prices: Vec<i32>) -> i32 {
let n = prices.len();
if n == 1 {
return 0;
}
let ku = k as usize;
let numRows = 2 * ku;
let mut dp: Vec<Vec<i32>> = Vec::with_capacity(numRows);
dp.resize(numRows, Vec::with_capacity(n));
for i in 0..numRows {
dp[i].resize(n, 0);
}
// Final Transaction
dp[0][n - 1] = prices[n - 1];
for i in (0..=(n-2)).rev() {
dp[0][i] = std::cmp::max(prices[i], dp[0][i+1]);
}
dp[1][n - 1] = 0;
for i in (0..=(n-2)).rev() {
dp[1][i] = std::cmp::max(dp[0][i+1] - prices[i], dp[1][i+1]);
}
// All other transactions
for j in 1..ku {
let boughtRow = 2 * j;
let preBoughtRow = boughtRow + 1;
dp[boughtRow][n - 1] = prices[n - 1];
dp[preBoughtRow][n - 1] = 0;
for i in (0..=(n-2)).rev() {
dp[boughtRow][i] = std::cmp::max(dp[boughtRow - 1][i+1] + prices[i], dp[boughtRow][i + 1]);
}
for i in (0..=(n-2)).rev() {
dp[preBoughtRow][i] = std::cmp::max(dp[boughtRow][i+1] - prices[i], dp[preBoughtRow][i + 1])
}
}
return dp[numRows - 1][0];
}Haskell Solution
As we saw in our first DP problem, we often don’t need as much memory as it initially seems. We filled out the “whole grid” for Rust, which helps make the algorithm more clear. But our Haskell solution will reflect the fact that we only actually need to pass along one preceding row (the pre-bought row) each time we loop through a transaction.
Let’s start by defining our edge case, as well as a few useful terms. We’ll define our indices in left-to-right order, but in all cases we’ll loop through them in reverse with foldr:
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else ...
where
n = V.length nums
lastPrice = nums V.! (n - 1)
idxs = ([0..(n-2)] :: [Int])
...Now we’ll define three different “loop” functions, all with the same pattern. We’ll use an IntMap Int to represent each “row” in our grid. So these functions will modify the IntMap for the row as we go along, while taking the new “index” we are populating. Let’s start with the base case, the first “bought” row, corresponding to our final transaction.
It will give us two options: sell or keep, following our algorithm. We insert the max of these into the map.
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else ...
where
n = V.length nums
lastPrice = nums V.! (n - 1)
idxs = ([0..(n-2)] :: [Int])
ibFold :: Int -> IM.IntMap Int -> IM.IntMap Int
ibFold i mp =
let sell = nums V.! i
keep = mp IM.! (i + 1)
in IM.insert i (max sell keep) mp
initialBought = foldr ibFold (IM.singleton (n-1) lastPrice) idxs
...We construct our initialBought row by folding, starting with a singleton of the last column base case.
Now we’ll write a function that, given a “bought” row, can construct the preceding “pre-bought” row. This will apply the “buy” and “stay” ideas in our algorithm and select between them. Choosing the “buy” option requires looking into the preceding “bought” row, while “stay” looking into a later index of the existing map:
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else ...
where
...
initialBought = foldr ibFold (IM.singleton (n-1) lastPrice) idxs
preBoughtFold :: IM.IntMap Int -> Int -> IM.IntMap Int -> IM.IntMap Int
preBoughtFold bought i preBought =
let buy = bought IM.! (i+1) - nums V.! i
stay = preBought IM.! (i+1)
in IM.insert i (max buy stay) preBought
initialPreBought = foldr (preBoughtFold initialBought) (IM.singleton (n-1) 0) idxsWe construct the initialPreBought row by applying this function with initialBought as the input. But we’ll use this for the rest of our “pre-bought” rows as well! First though, we need a more general loop for the rest of our “bought” rows.
This function has the same structure as pre-bought, just applying the “sell” and “keep” rules instead of “buy” and “stay”:
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else ...
where
...
boughtFold :: IM.IntMap Int -> Int -> IM.IntMap Int -> IM.IntMap Int
boughtFold preBought i bought =
let sell = preBought IM.! (i+1) + nums V.! i
keep = bought IM.! (i+1)
in IM.insert i (max sell keep) boughtNow we’re ready for our core loop! This will loop through every transaction except the base case. It takes only the preceding “pre-bought” row and the transaction counter. Once the counter reaches k, we return the first value in this row. Otherwise, we run the “bought” loop to produce a next “bought” row, and we pass this in to the “pre-bought” loop to produce a new “pre-bought” row. This becomes the input to our recursive call:
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else loop 1 initialPreBought
where
...
loop :: Int -> IM.IntMap Int -> Int
loop i preBought = if i >= k then preBought IM.! 0
else
let bought' = foldr (boughtFold preBought) (IM.singleton (n-1) lastPrice) idxs
preBought' = foldr (preBoughtFold bought') (IM.singleton (n-1) 0) idxs
in loop (i + 1) preBought'As you can see above, we complete the solution by calling our loop with the initial “pre-bought” row, and a transaction counter of 1!
Here’s our full Haskell solution:
maxProfit :: V.Vector Int -> Int -> Int
maxProfit nums k = if n == 1 then 0
else loop 1 initialPreBought
where
n = V.length nums
lastPrice = nums V.! (n - 1)
idxs = ([0..(n-2)] :: [Int])
ibFold :: Int -> IM.IntMap Int -> IM.IntMap Int
ibFold i mp =
let sell = nums V.! i
keep = mp IM.! (i + 1)
in IM.insert i (max sell keep) mp
initialBought = foldr ibFold (IM.singleton (n-1) lastPrice) idxs
preBoughtFold :: IM.IntMap Int -> Int -> IM.IntMap Int -> IM.IntMap Int
preBoughtFold bought i preBought =
let buy = bought IM.! (i+1) - nums V.! i
stay = preBought IM.! (i+1)
in IM.insert i (max buy stay) preBought
initialPreBought = foldr (preBoughtFold initialBought) (IM.singleton (n-1) 0) idxs
boughtFold :: IM.IntMap Int -> Int -> IM.IntMap Int -> IM.IntMap Int
boughtFold preBought i bought =
let sell = preBought IM.! (i+1) + nums V.! i
keep = bought IM.! (i+1)
in IM.insert i (max sell keep) bought
loop :: Int -> IM.IntMap Int -> Int
loop i preBought = if i >= k then preBought IM.! 0
else
let bought' = foldr (boughtFold preBought) (IM.singleton (n-1) lastPrice) idxs
preBought' = foldr (preBoughtFold bought') (IM.singleton (n-1) 0) idxs
in loop (i + 1) preBought'Conclusion
That’s the last LeetCode solution we’re going to write for now! Hopefully you’ve got a good impression now on the differences between dynamic programming in Haskell when compared to a language like Rust. To learn more about the basics of these Haskell solutions, take a look at our course, Solve.hs! You’ll also get a ton of practice with hundreds of exercise problems in the course!
Next week we will switch gears, and start working on some interesting parsing problems!