State-machine DP (stocks)

44 s · title card, the worked example step by step, then a recap. The captions carry the explanation; there is no narration. Download the MP4.

01The intuition

Start with what you would do by hand

Prices on five days are [1, 2, 3, 0, 2]. You may buy one share, sell it later, and buy again, as many times as you like, but after a sale you must rest one day before buying (a cooldown). What is the most money you can make?

By hand you would try plans. Buy at 1, sell at 3, rest on day 3, buy at 0, sell at 2: profit 2 + 2 = 4? No, check the cooldown: you sold on day 2, so day 3 is the rest day and you cannot buy at 0. Buy at 1, sell at 2, rest, buy at 0, sell at 2: profit 1 + 2 = 3. There are many plans, and each day you have two or three choices, so for 30 days there are billions of plans. The by-hand method also makes it easy to break a rule without noticing, as the first attempt did.

Build the structure, one piece at a time

Instead of remembering plans, remember only what situation you are in at the end of each day. There are three situations, called states:

• hold: you own a share right now.
• sold: you sold today, so tomorrow is a forced rest.
• rest: you own nothing and are free to buy.

For each state keep one number: the most cash you could have while being in that state. Each new day, each state can be reached from a few earlier states, and you take the best:

• hold today = max(hold yesterday, rest yesterday − price): keep holding, or buy today from a free day.
• sold today = hold yesterday + price: you must have held a share to sell it.
• rest today = max(rest yesterday, sold yesterday): stay free, or come off cooldown.

Before day 0, rest = 0, hold = −∞ (impossible), sold = −∞. Now run the days:

Day 3 shows the cooldown working: hold becomes rest − price = 1 − 0 = 1, and that rest value, 1, came from a sale on day 1, not from the sale on day 2. The answer is the best of sold and rest on the last day: max(3, 2) = 3, which is buy at 1, sell at 2, rest, buy at 0, sell at 2.

Why it is correct

The invariant: after day d, hold is the best cash over all legal plans that end day d owning a share, and likewise for sold and rest. A plan that ends day d holding a share either held it yesterday too or bought it today; buying today is only legal from a free day, which is exactly the rest state, never the sold state. So the three update rules list every legal way into each state and nothing illegal. Taking the maximum over legal ways keeps each number the true best, so the invariant passes from one day to the next.

hold_d = max(hold_d−1, rest_d−1 − p_d), sold_d = hold_d−1 + p_d, rest_d = max(rest_d−1, sold_d−1)

Why it is fast

Each day does three updates, each a maximum of two numbers. For n days that is about 6n operations: O(n) time. Only yesterday's three numbers are needed, so memory is O(1). The by-hand search grew like 2ⁿ. With a limit of k transactions the states become (holding or not) × (transactions used), giving O(n · k).

How to recognise it in an interview

Cues: a sequence of days or steps, a small set of situations you can be in (holding, resting, used k of something), and rules about which situation can follow which. Draw the states as circles and the legal moves as arrows, then write one line per circle. Best Time to Buy and Sell Stock with Cooldown is the three-circle version above. Best Time to Buy and Sell Stock II drops the sold circle. Best Time to Buy and Sell Stock IV copies the circles once per allowed transaction.

02The idea in one line

States: holding / not holding (plus cooldown); transitions carry the best cash.

Reach for it when: Stock II/III/IV, with cooldown, with fee.

Cost: O(n · k) time.

03The walk-through, step by step

The video works one example: Stock with Cooldown on [1,2,3,0,2]. These are its captions in order, so you can follow without the video or check a step you missed.

1. Three states: hold a share, just sold (cooldown), or rest
2. Day 0, price 1: hold = -1, sold = −∞, rest = 0
3. Day 1, price 2: hold = -1, sold = 1, rest = 0
4. Day 2, price 3: hold = -1, sold = 2, rest = 1
5. Day 3, price 0: hold = 1, sold = -1, rest = 2
6. Day 4, price 2: hold = 1, sold = 3, rest = 2
7. Why rest − p for buying: you can only buy from rest, never right after a sale
8. max(sold, rest) at the end = 3: buy 1 sell 3 · cool · buy 0 sell 2
9. One pass, three numbers of state: O(n) time, O(1) space

04Asked at Google

6 Google-tagged LeetCode problems use this pattern. Frequency is LeetCode's Google-tag score (higher means reported more often); "last 6 months" marks problems reported in the most recent window of the export.

05Related patterns

Other patterns in Dynamic Programming: 1-D linear DP, Kadane (best ending here), Longest increasing subsequence, Grid path DP, 0/1 knapsack (subset sum), Unbounded knapsack (coin change), Two-sequence DP (LCS / edit distance), Wildcard / regex matching DP, Palindromes: expand or table, Prefix segmentation DP, Interval DP (last operation), Bitmask DP over subsets, Memoised DFS (top-down), Digit DP, Game DP (minimax on a range), Counting ways with modulo, DP over sorted keys (chains).