Best Time to Buy And Sell Stock With Cooldown

The problem

prices lists the price of one share on each day. You may buy and sell as often as you like, but you can own at most one share at a time, and after a sale you must wait one full day before you buy again. Each trade earns its selling price minus its buying price. Return the largest total profit you can make. If no trade makes money, return 0.

What to write

function solve(prices: number[]): number

Input
prices = [1,2,3,0,2]

Output
3

Input
prices = [1]

Output
0

Constraints

Prices are nonnegative.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   You never have to trade, so the answer is at least 0. Without the waiting day you could collect every rise in price. With it, a sale today can block a cheap buy tomorrow, so the best move today depends on prices you have not looked at yet.

2. Hint 2: Ask a smaller question

   Stand at the end of some day in 1, 2, 3, 0, 2. What do you need to know about your past to decide what you may do tomorrow? Is the full list of your trades needed, or only a few facts about where you stand right now?

3. Hint 3: Name the pattern

   What you may do tomorrow depends only on how today ends: holding a share, sold today, or resting, which means no share and no sale today. For every day and each of these three situations, keep the most money you could have. Days in one direction and situations in the other make a small table: two-dimensional dynamic programming. Read the lesson on two-dimensional dynamic programming.

4. Hint 4: State the invariant

   After each day keep three numbers: the most money you can have if the day ends holding a share, if it ends with a sale made that day, and if it ends resting, with no share and no sale that day. Money starts at 0; a buy subtracts the price, a sale adds it. Each number needs only yesterday's three.

5. Hint 5: Sketch the pseudocode

   holding = sold = impossible (minus infinity), resting = 0. For each price: new sold = holding + price; new holding = larger of holding and (resting - price); new resting = larger of resting and sold. Replace all three together, so every right-hand side uses yesterday's values. Return the larger of resting and sold.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Track holding a stock, selling today, and resting. A buy can only follow a resting state, which enforces the cooldown.

The slow way

Try every schedule: on each day buy, sell or wait, wherever the rules allow, and keep the best result. It is correct, but the number of schedules grows fast. The 5 days of example 1 allow 21 schedules, and 30 days allow 26931732.

The insight

What you may do tomorrow depends only on how today ends: holding a share, having sold today, or resting, which means no share and no sale today. The trades that led there do not matter. So each of the three endings needs one number, the most money so far, and tomorrow's three follow from today's three. Carrying best answers forward like this is dynamic programming.

Step by step

1. Start before day one: resting has 0, holding and sold are impossible. Each value is the most money you can have in that situation. Before any day you own nothing and have sold nothing, so only resting is possible. Minus infinity marks the other two as impossible: it loses every comparison.
2. For each day's price, set aside yesterday's sold value first. The sold value is about to be overwritten with today's sale, but the resting update further down still needs yesterday's. Without the saved copy, resting would include a sale made today.
3. Sold today: yesterday's holding value plus today's price. You can only sell a share you already held last night, and the sale brings in today's price. While holding is still impossible, this sum stays impossible too.
4. Holding: the better of keeping the share and buying after a rest. You hold tonight if you held last night and kept the share, or if you were resting last night and buy today, which costs today's price. A buy never starts from the sold value: that is the waiting day.
5. Resting: the better of resting yesterday and having sold yesterday. You end today with no share and no sale if yesterday ended the same way, or if you sold yesterday and today is your waiting day. This uses the saved copy, not the sold value written a moment ago.
6. Return the larger of the resting value and the sold value. After the last day the best schedule owns no share, so it ends either resting or with a sale on that last day. A share still held was paid for and never sold, so holding is never better than resting.

The same numbered steps are the comments in each solution.

function solve(prices) {
  // 1. Start before day one: resting has 0, holding and sold are impossible
  let hold = -Infinity,
    sold = -Infinity,
    rest = 0;
  // 2. For each day's price, set aside yesterday's sold value first
  for (const price of prices) {
    const oldSold = sold;
    // 3. Sold today: yesterday's holding value plus today's price
    sold = hold + price;
    // 4. Holding: the better of keeping the share and buying after a rest
    hold = Math.max(hold, rest - price);
    // 5. Resting: the better of resting yesterday and having sold yesterday
    rest = Math.max(rest, oldSold);
  }
  // 6. Return the larger of the resting value and the sold value
  return Math.max(rest, sold);
}

function solve(prices: number[]): number {
  // 1. Start before day one: resting has 0, holding and sold are impossible
  let hold = -Infinity,
    sold = -Infinity,
    rest = 0;
  // 2. For each day's price, set aside yesterday's sold value first
  for (const price of prices) {
    const oldSold = sold;
    // 3. Sold today: yesterday's holding value plus today's price
    sold = hold + price;
    // 4. Holding: the better of keeping the share and buying after a rest
    hold = Math.max(hold, rest - price);
    // 5. Resting: the better of resting yesterday and having sold yesterday
    rest = Math.max(rest, oldSold);
  }
  // 6. Return the larger of the resting value and the sold value
  return Math.max(rest, sold);
}

def solve(prices):
    # 1. Start before day one: resting has 0, holding and sold are impossible
    hold = sold = float("-inf")
    rest = 0
    # 2. For each day's price, set aside yesterday's sold value first
    for price in prices:
        old_sold = sold
        # 3. Sold today: yesterday's holding value plus today's price
        sold = hold + price
        # 4. Holding: the better of keeping the share and buying after a rest
        hold = max(hold, rest - price)
        # 5. Resting: the better of resting yesterday and having sold yesterday
        rest = max(rest, old_sold)
    # 6. Return the larger of the resting value and the sold value
    return max(rest, sold)

Example 1, traced

Input
prices = [1,2,3,0,2]

Output
3

Why it works

Each value is the most money over all schedules that end the day in its situation. A schedule that sold today held a share yesterday. One that holds either held yesterday or rested yesterday and bought today. One that rests either rested or sold yesterday. These cases cover every allowed schedule and no forbidden one, since a buy never follows a sale directly.

Time and space

Time O(n), extra space O(1). Each price is handled once with a fixed handful of additions and comparisons, so the time is O(n) for n prices. A full table would have a row per day and a column per situation, but only the latest row, three numbers and one saved copy, is kept: O(1) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Update resting from yesterday's sold value, not today's. Without the saved copy, a sale today counts as rest at once and you can buy the next morning: the prices 1, 2, 3, 0, 2 give 4 instead of 3.
• Start holding as impossible, not as 0. A 0 means you own a share that cost nothing, and that share gets sold on day one: the single price 1 returns 1 instead of 0.
• Return the larger of resting and sold, not resting alone. The best schedule often sells on the last day, and that sale has not reached the resting value yet: 1, 2, 3, 0, 2 would give 2.

Pattern

Two-dimensional dynamic programming. Builds on Best Time to Buy And Sell Stock and House Robber.

Keep going