Buying and selling, with a state machine

A family that looks like a dozen separate problems and is one idea: track what state you are in, and the transitions between them.

Two states

At the end of each day you either hold a share or you do not.

def max_profit(prices):
    hold, free = float('-inf'), 0
    for p in prices:
        hold = max(hold, free - p)       # keep holding, or buy today
        free = max(free, hold + p)       # stay out, or sell today
    return free

hold is the best balance while holding a share, free the best while not. The answer is free, because ending while holding is never better.

This allows unlimited transactions. Two variables, one pass, O(1) space.

Adding constraints

Each variation adds states or dimensions rather than a new idea:

At most one transaction: hold = max(hold, -p). Buying resets to zero rather than using accumulated profit, so only one purchase can ever be made.

At most k transactions: two arrays of length k+1.

def max_profit_k(prices, k):
    hold = [float('-inf')] * (k + 1)
    free = [0] * (k + 1)
    for p in prices:
        for t in range(1, k + 1):
            hold[t] = max(hold[t], free[t-1] - p)   # this purchase starts transaction t
            free[t] = max(free[t], hold[t] + p)
    return free[k]

With a fee: subtract it on sale, free = max(free, hold + p - fee).

With a cooldown: a third state for the day after selling, because you may not buy immediately.

hold, sold, rest = float('-inf'), 0, 0
for p in prices:
    prev_sold = sold
    sold = hold + p
    hold = max(hold, rest - p)
    rest = max(rest, prev_sold)

Why this generalises

The pattern is a state machine: a small set of situations and the legal moves between them. Drawing the states and arrows first makes the code fall out, and it is the recognition worth taking from the family.

     buy
free ---> hold
  ^         |
  |  sell   |
  +---------+

Add cooldown and a node appears between sell and buy. Add transaction limits and each node gains a counter dimension.

The order of updates

In the unlimited version, hold is updated before free, so free uses today's hold. That allows buying and selling on the same day, which nets zero and changes nothing. In the cooldown version the order matters, which is why prev_sold is saved first.

This is the same staleness question as a rolling array, appearing as variables rather than as an array index.