Best Time to Buy and Sell Stock: DP as a State Machine

"Best time to buy and sell stock" is not one problem but a family: one trade, as many trades as you like, a fee per trade, a cooldown after selling, at most k trades. Each variant has its own trick if you learn them separately. Seen as a small state machine, they are all the same loop: a couple of numbers, one per state, updated once per day. This article builds that loop for four variants and checks every one against a brute force that tries every sequence of actions.

The problem it solves

You get a list of daily prices. You may buy one share, sell it later, and repeat, holding at most one share at a time. Maximise the profit. The variants change the rules:

• One transaction: buy once, sell once.
• Unlimited transactions: as many buy-sell pairs as you like.
• Transaction fee: each completed sale costs a fixed fee.
• Cooldown: after a sale you must skip one day before buying again.

Enumerating every sequence of buy, sell and wait is exponential. The state machine is O(n) time and O(1) space for each of these variants.

The intuition

At the end of any day you are in one of two situations: holding a share, or free of one. For each, keep the best balance you could have in that situation, counting purchases as negative cash:

• hold: the best balance at the end of the day while owning a share.
• free: the best balance while owning none.

A day's move is a transition between states. You can stay where you are, or cross over:

• new hold = max(hold, free - price): keep holding, or buy today.
• new free = max(free, hold + price - fee): stay in cash, or sell today.

Both new values come from yesterday's pair, which is why the update is written as one simultaneous assignment. The answer is free at the end: finishing with an unsold share never helps.

Each variant is a small change to the machine:

• Unlimited: the rule above with a fee of 0.
• One transaction: buying may only happen from the starting cash of 0, so the buy move becomes -price instead of free - price.
• Cooldown: split free into two states, sold (sold today, so tomorrow is blocked) and rest (free and allowed to buy). Buying comes only from rest; sold becomes rest the next day.

Watch it run

The animation draws the two states for prices [7, 1, 5, 3, 6, 4] with a fee of 2. Day one: holding is -7, free is 0. Price 1: buying today would leave -1, selling would leave -8, and each state keeps the better of yesterday and the move, so hold becomes -1 and free stays 0. Price 5: buying would leave -5, selling would leave 2, so free becomes 2. Price 3: buying would leave -1, which only ties hold, and selling would leave 0, so nothing changes. Price 6: selling would leave 3, the new best. Price 4: selling would leave only 1. The last frame reads 3: bought at 1, sold at 6, minus the fee of 2. No trade list was ever stored; two numbers were enough.

The code

Four variants, one shape. Starting hold at minus infinity means "no share yet", so the first day is handled by the same line as every other day:

NEG = float("-inf")

def with_fee(prices, fee=0):
    hold, free = NEG, 0
    for p in prices:
        hold, free = max(hold, free - p), max(free, hold + p - fee)
    return free

def one_transaction(prices):
    hold, free = NEG, 0
    for p in prices:
        hold, free = max(hold, -p), max(free, hold + p)   # buy only from the initial 0
    return free

def with_cooldown(prices):
    hold, sold, rest = NEG, NEG, 0
    for p in prices:
        hold, sold, rest = max(hold, rest - p), hold + p, max(rest, sold)
    return max(sold, rest)

prices = [7, 1, 5, 3, 6, 4]
print(with_fee(prices, 2))                                  # 3
print(one_transaction(prices), with_fee(prices))            # 5 7
print(with_cooldown([1, 2, 3, 0, 2]))                       # 3
print(with_fee([1, 3, 2, 8, 4, 9], 2), with_fee([], 2))     # 8 0

Why the states are assigned on one line: the same cooldown machine updated one state at a time lets a sale and the next purchase share a day:

def with_cooldown_in_sequence(prices):
    hold, sold, rest = NEG, NEG, 0
    for p in prices:
        hold = max(hold, rest - p)
        sold = hold + p                          # reads today's hold
        rest = max(rest, sold)                   # reads today's sold: no cooldown
    return max(sold, rest)

print(with_cooldown_in_sequence([5, 8, 6, 8]), with_cooldown([5, 8, 6, 8]))   # 5 3

A brute force that tries every sequence of wait, buy and sell, day by day, with no memoisation, against all four on 1,500 short random price lists:

import random

def brute(prices, fee=0, max_trades=None, cooldown=False):
    def go(day, holding, blocked, trades, cash):
        if day == len(prices):
            return cash if not holding else NEG
        p = prices[day]
        best = go(day + 1, holding, False, trades, cash)             # wait
        if holding:
            best = max(best, go(day + 1, False, cooldown, trades, cash + p - fee))
        elif not blocked and (max_trades is None or trades < max_trades):
            best = max(best, go(day + 1, True, False, trades + 1, cash - p))
        return best
    return go(0, False, False, 0, 0)

random.seed(17)
ok = True
for _ in range(1500):
    ps = [random.randint(1, 9) for _ in range(random.randint(0, 9))]
    fee = random.randint(0, 3)
    ok &= with_fee(ps, fee) == brute(ps, fee)
    ok &= one_transaction(ps) == brute(ps, max_trades=1)
    ok &= with_cooldown(ps) == brute(ps, cooldown=True)
print(ok)                                                   # True

The complexity

• Brute force: O(3^n) sequences in the worst case.
• State machine: O(n) time for every variant here, and O(1) space: two or three numbers.
• At most k transactions adds a counter to the state: hold[j] and free[j] for j trades used, O(nk) time and O(k) space. Same machine, more states.

Where it goes wrong

• Updating the states one at a time. In the two-state machine it happens to be harmless, because buying and selling on the same day nets zero or loses the fee. In the cooldown machine it is not: sold reads today's hold and rest reads today's sold, so a sale and the next purchase can share a day. On [5, 8, 6, 8] that reports 5 instead of 3. Assign all states at once.
• Charging the fee twice. Charge it once per completed trade, on the buy or on the sell, not both.
• Starting hold at 0. That pretends you own a free share. Start at minus infinity, or at -prices[0] as the lesson does.
• Returning max(hold, free). A held share at the end is money spent, not earned.
• Buying the day after a sale in the cooldown variant. Buying must come from rest, never from sold.

When it shows up in interviews

The one-transaction version is a very common easy question, usually solved with a running minimum. The follow-ups are where this pattern pays: "now unlimited trades", "now with a fee", "now with a cooldown", "now at most two trades". Candidates who memorised four tricks get stuck; candidates who draw the states and their arrows can answer each follow-up in a minute. The same idea, a few named states updated per step, also solves house robber and parsing problems written as finite automata.

How to say it in an interview

"At the end of each day I am either holding a share or free. I keep the best balance for each: hold is the max of keeping or buying today from free, free is the max of staying in cash or selling today from hold, minus the fee. Both update from yesterday's values at once, and the answer is free at the end. It is O(n) time and O(1) space. For a cooldown I split free into just-sold and resting, and for at most k trades I index the states by trades used."

The two-number idea first appears in house robber, and the one-transaction version is a cousin of Kadane's algorithm.