Solution to Best Time to Buy and Sell Stock with Cooldown

This is my solution for the LeetCode problem number 309, Best Time to Buy and Sell Stock with Cooldown.

This is a quite simple problem which can be addressed in O(1) space and O(n) time using dynamic programming. However, the O(n) space solution seems easier to arrive at. The subproblem explored through dynamic programming is that starting day \(i\) with no stock and after the cooldown has ended is equivalent to solving the original problem with only days \(i\) through \(n\).

In this solution, 0 is used to indicate that no stock is owned, 1 is used to indicate that stock is owned and 2 is used to indicate cooldown. The algorithm keeps three values for each day: the maximum profit by starting without stock, the maximum profit by starting with stock and the maximum profit by starting under cooldown. After the bottom-up solving is done, the answer to the problem question is what is the maximum profit for starting at the first day without stock.

int maxProfit(const vector<int> &prices) {
  vector<array<int, 3>> m(prices.size() + 1);
  for (int i = prices.size() - 1; i >= 0; i--) {
    m[i][0] = max(m[i + 1][0], m[i + 1][1] - prices[i]);
    m[i][1] = max(m[i + 1][1], m[i + 1][2] + prices[i]);
    m[i][2] = max(m[i + 1][2], m[i + 1][0]);
  }
  return m[0][0];
}

Now, because evaluating day \(i\) only requires information about day \(i + 1\), the vector of arrays can become just two arrays, as follows.

int maxProfit(const vector<int> &prices) {
  array<int, 3> d1{};
  array<int, 3> d2{};
  for (int i = prices.size() - 1; i >= 0; i--) {
    d2 = d1;
    d1[0] = max(d2[0], d2[1] - prices[i]);
    d1[1] = max(d2[1], d2[2] + prices[i]);
    d1[2] = max(d2[2], d2[0]);
  }
  return d1[0];
}