Skip to content
Codeloom
DSA

DP State Machine: Buy and Sell Stock Mastery

Model dynamic programming as state machines — solve all Buy and Sell Stock variants (I-IV, cooldown, fee), understand state transitions, and build a general framework with Python.

·11 min read · By Codeloom
Intermediate 18 min read

What you'll learn

  • How to model DP problems as state machines with explicit transitions
  • All Buy and Sell Stock variants from LeetCode (I through IV)
  • How cooldown and transaction fee modify the state machine
  • A general framework that solves ANY stock trading problem
  • State reduction techniques for O(1) space solutions
  • How to apply state machine thinking beyond stock problems

Prerequisites

Many DP problems become clearer when you think of them as state machines. At each step, you are in some state and can transition to another state by taking an action. The DP value at each state represents the optimal outcome so far. The Buy and Sell Stock family of problems is the perfect case study for this approach.

DP State Machine — state diagram with transitions for Buy and Sell Stock with Cooldown


1. The State Machine Framework

In a state machine DP:

  1. Define the states: What are the possible situations you can be in?
  2. Define transitions: From each state, what actions can you take?
  3. Define the DP value: What are you optimising at each state?
  4. Base cases: What is the initial state?
  5. Answer: Which state(s) contain the final answer?

For stock problems, the states revolve around:

  • Whether you currently hold a stock
  • How many transactions you have completed
  • Whether you are in a cooldown period

2. Best Time to Buy and Sell Stock I (One Transaction)

Problem: You can make at most one transaction (buy one day, sell another later day). Maximise profit.

This does not even need state machine DP — a single pass suffices. But let us frame it as a state machine to build intuition.

States: NO_STOCK (never bought or already sold), HOLDING (bought but not sold).

def max_profit_one(prices):
    """
    LeetCode 121: Best Time to Buy and Sell Stock.

    States: no_stock, holding
    At most 1 buy and 1 sell.

    Time: O(n), Space: O(1)
    """
    no_stock = 0          # max profit when not holding
    holding = float('-inf')  # max profit when holding (haven't bought yet)

    for price in prices:
        # If we sell today: holding -> no_stock
        no_stock = max(no_stock, holding + price)
        # If we buy today: no_stock(before) -> holding
        # Since only 1 transaction, buying costs us -price from initial 0
        holding = max(holding, -price)

    return no_stock


print(max_profit_one([7, 1, 5, 3, 6, 4]))  # 5 (buy@1, sell@6)
print(max_profit_one([7, 6, 4, 3, 1]))     # 0 (no transaction)

The classic single-pass solution (track min price) is equivalent:

def max_profit_one_classic(prices):
    """Classic approach: track minimum price seen so far."""
    min_price = float('inf')
    max_profit = 0

    for price in prices:
        min_price = min(min_price, price)
        max_profit = max(max_profit, price - min_price)

    return max_profit

3. Best Time to Buy and Sell Stock II (Unlimited Transactions)

Problem: You can make as many transactions as you want (but must sell before buying again).

States: no_stock, holding.

def max_profit_unlimited(prices):
    """
    LeetCode 122: Best Time to Buy and Sell Stock II.

    Unlimited transactions.

    Time: O(n), Space: O(1)
    """
    no_stock = 0
    holding = float('-inf')

    for price in prices:
        prev_no_stock = no_stock

        # Sell: holding -> no_stock
        no_stock = max(no_stock, holding + price)
        # Buy: no_stock -> holding (can buy again after selling)
        holding = max(holding, prev_no_stock - price)

    return no_stock


print(max_profit_unlimited([7, 1, 5, 3, 6, 4]))  # 7 (1->5: +4, 3->6: +3)
print(max_profit_unlimited([1, 2, 3, 4, 5]))     # 4 (buy@1, sell@5)

The greedy equivalent: add up all positive differences.

def max_profit_unlimited_greedy(prices):
    """Greedy: collect every upward move."""
    return sum(max(0, prices[i] - prices[i - 1]) for i in range(1, len(prices)))

4. Best Time to Buy and Sell Stock III (At Most 2 Transactions)

Problem: at most 2 transactions.

States: for each transaction k in {1, 2}, track holding[k] and no_stock[k].

def max_profit_two(prices):
    """
    LeetCode 123: Best Time to Buy and Sell Stock III.

    At most 2 transactions.

    Time: O(n), Space: O(1)
    """
    # State: (transactions_completed, holding_or_not)
    # no_stock_1: sold after 1st transaction (or haven't started)
    # holding_1: bought for 1st transaction
    # no_stock_2: sold after 2nd transaction
    # holding_2: bought for 2nd transaction

    hold1 = float('-inf')
    sell1 = 0
    hold2 = float('-inf')
    sell2 = 0

    for price in prices:
        # Order matters: process sell2 -> hold2 -> sell1 -> hold1
        sell2 = max(sell2, hold2 + price)
        hold2 = max(hold2, sell1 - price)
        sell1 = max(sell1, hold1 + price)
        hold1 = max(hold1, -price)

    return sell2


print(max_profit_two([3, 3, 5, 0, 0, 3, 1, 4]))  # 6 (0->3: +3, 1->4: +3)
print(max_profit_two([1, 2, 3, 4, 5]))            # 4 (one transaction suffices)

5. Best Time to Buy and Sell Stock IV (At Most K Transactions)

Problem: at most k transactions.

def max_profit_k(k, prices):
    """
    LeetCode 188: Best Time to Buy and Sell Stock IV.

    At most k transactions.

    Time: O(n * k), Space: O(k)
    """
    n = len(prices)
    if not prices or k == 0:
        return 0

    # If k >= n//2, it's equivalent to unlimited transactions
    if k >= n // 2:
        return sum(max(0, prices[i] - prices[i - 1]) for i in range(1, n))

    # hold[j] = max profit when holding stock after j-th buy
    # sell[j] = max profit after j-th sell
    hold = [float('-inf')] * (k + 1)
    sell = [0] * (k + 1)

    for price in prices:
        for j in range(1, k + 1):
            sell[j] = max(sell[j], hold[j] + price)
            hold[j] = max(hold[j], sell[j - 1] - price)

    return sell[k]


print(max_profit_k(2, [3, 2, 6, 5, 0, 3]))  # 7 (2->6: +4, 0->3: +3)
print(max_profit_k(1, [2, 4, 1]))            # 2

6. Best Time to Buy and Sell Stock with Cooldown

Problem: unlimited transactions, but after selling you must wait one day before buying again.

States: rest (no stock, can buy), hold (holding stock), cool (just sold, must wait).

def max_profit_cooldown(prices):
    """
    LeetCode 309: Best Time to Buy and Sell Stock with Cooldown.

    After selling, must wait 1 day before buying.

    State transitions:
      rest -> rest (skip)
      rest -> hold (buy)
      hold -> hold (keep)
      hold -> cool (sell)
      cool -> rest (mandatory wait)

    Time: O(n), Space: O(1)
    """
    rest = 0                # can buy or skip
    hold = float('-inf')    # holding stock
    cool = float('-inf')    # just sold, in cooldown

    for price in prices:
        prev_rest = rest
        prev_hold = hold
        prev_cool = cool

        rest = max(prev_rest, prev_cool)    # skip or exit cooldown
        hold = max(prev_hold, prev_rest - price)  # keep or buy
        cool = prev_hold + price             # sell

    return max(rest, cool)


print(max_profit_cooldown([1, 2, 3, 0, 2]))  # 3 (buy@1,sell@3,cool,buy@0,sell@2)
print(max_profit_cooldown([1]))               # 0

7. Best Time to Buy and Sell Stock with Transaction Fee

Problem: unlimited transactions, but each transaction incurs a fee.

def max_profit_with_fee(prices, fee):
    """
    LeetCode 714: Best Time to Buy and Sell Stock with Transaction Fee.

    Each complete transaction (buy + sell) costs `fee`.

    Time: O(n), Space: O(1)
    """
    no_stock = 0
    holding = float('-inf')

    for price in prices:
        prev_no_stock = no_stock

        # Sell: subtract fee
        no_stock = max(no_stock, holding + price - fee)
        # Buy
        holding = max(holding, prev_no_stock - price)

    return no_stock


print(max_profit_with_fee([1, 3, 2, 8, 4, 9], 2))  # 8
# buy@1, sell@8 (-2 fee) = +5, buy@4, sell@9 (-2 fee) = +3 -> total 8
print(max_profit_with_fee([1, 3, 7, 5, 10, 3], 3))  # 6

8. The General Stock Trading Framework

All stock problems follow the same pattern. Here is a unified solution:

def max_profit_general(
    prices,
    max_transactions=float('inf'),
    cooldown_days=0,
    fee=0
):
    """
    General framework for stock trading problems.

    Args:
        prices: list of daily prices
        max_transactions: max number of buy-sell pairs (inf = unlimited)
        cooldown_days: number of days to wait after selling (0 = none)
        fee: fee per complete transaction

    Returns:
        Maximum profit.
    """
    n = len(prices)
    if n <= 1:
        return 0

    # Determine effective k
    if max_transactions >= n // 2:
        # Unlimited transactions — use O(n) solution
        no_stock = 0
        holding = float('-inf')
        # For cooldown, keep a history of no_stock values
        no_stock_history = [0] * (cooldown_days + 1)

        for price in prices:
            prev_holding = holding

            # Sell
            no_stock = max(no_stock, holding + price - fee)
            # Buy (use no_stock from cooldown_days ago)
            holding = max(holding, no_stock_history[0] - price)

            # Shift history
            no_stock_history = [no_stock] + no_stock_history[:-1]

        return max(no_stock, 0)

    # Limited transactions — use O(n*k) solution
    k = max_transactions
    hold = [float('-inf')] * (k + 1)
    sell = [0] * (k + 1)

    for price in prices:
        for j in range(1, k + 1):
            sell[j] = max(sell[j], hold[j] + price - fee)
            hold[j] = max(hold[j], sell[j - 1] - price)

    return sell[k]


# Test all variants
prices = [7, 1, 5, 3, 6, 4]
print("Stock I:", max_profit_general(prices, max_transactions=1))         # 5
print("Stock II:", max_profit_general(prices))                             # 7
print("Stock III:", max_profit_general(prices, max_transactions=2))        # 7

prices2 = [1, 2, 3, 0, 2]
print("Cooldown:", max_profit_general(prices2, cooldown_days=1))           # 3

prices3 = [1, 3, 2, 8, 4, 9]
print("With fee:", max_profit_general(prices3, fee=2))                     # 8

9. State Machine DP Beyond Stocks

The state machine pattern applies to many other problems.

Example: Paint House

Problem: paint n houses with 3 colours such that no two adjacent houses have the same colour. Each colour has a different cost per house.

def min_cost_paint_houses(costs):
    """
    LeetCode 256: Paint House.

    States: R, G, B (colour of current house)
    Transition: each state can transition to the other two states.

    Time: O(n), Space: O(1)
    """
    if not costs:
        return 0

    r, g, b = costs[0]

    for i in range(1, len(costs)):
        new_r = costs[i][0] + min(g, b)
        new_g = costs[i][1] + min(r, b)
        new_b = costs[i][2] + min(r, g)
        r, g, b = new_r, new_g, new_b

    return min(r, g, b)


costs = [[17, 2, 17], [16, 16, 5], [14, 3, 19]]
print(min_cost_paint_houses(costs))  # 10 (green + blue + green: 2+5+3)

Example: Decode Ways with State Machine

def num_decodings(s):
    """
    LeetCode 91: Decode Ways.

    States represent whether we're mid-two-digit decode or not.
    """
    if not s or s[0] == '0':
        return 0

    n = len(s)
    # dp[i] = number of ways to decode s[0..i-1]
    prev2 = 1  # dp[i-2]
    prev1 = 1  # dp[i-1]

    for i in range(1, n):
        curr = 0

        # Single digit decode
        if s[i] != '0':
            curr += prev1

        # Two digit decode
        two_digit = int(s[i - 1:i + 1])
        if 10 <= two_digit <= 26:
            curr += prev2

        prev2, prev1 = prev1, curr

    return prev1


print(num_decodings("226"))  # 3 ("BZ", "VF", "BBF")
print(num_decodings("06"))   # 0

10. Visualising State Machines

The best way to approach state machine DP:

  1. Draw the states as circles.
  2. Draw transitions as arrows with conditions and costs.
  3. Write the recurrence from the diagram.
  4. Code it — each state becomes a variable.

For the stock with cooldown problem, the diagram is:

   +---------+      buy (-price)      +---------+
   |  REST   |  ------------------>   |  HOLD   |
   |no stock |  <------ wait ------   | holding |
   +---------+                        +---------+
       ^                                  |
       |              sell (+price)       |
       |          +----------+           |
       +--------- | COOLDOWN | <---------+
         wait     +----------+

11. Complexity Summary

ProblemStatesTimeSpace
Stock I (1 txn)2O(n)O(1)
Stock II (unlimited)2O(n)O(1)
Stock III (2 txns)4O(n)O(1)
Stock IV (k txns)2kO(nk)O(k)
Stock + cooldown3O(n)O(1)
Stock + fee2O(n)O(1)
Paint House (3 colours)3O(n)O(1)

12. Practice Problems

ProblemPlatformKey Technique
Best Time to Buy and Sell Stock (LC 121)LeetCode1 transaction
Best Time to Buy and Sell Stock II (LC 122)LeetCodeUnlimited
Best Time to Buy and Sell Stock III (LC 123)LeetCode2 transactions
Best Time to Buy and Sell Stock IV (LC 188)LeetCodeK transactions
Best Time with Cooldown (LC 309)LeetCodeCooldown state
Best Time with Fee (LC 714)LeetCodeFee on sell
Paint House (LC 256)LeetCode3-state machine
Paint House II (LC 265)LeetCodeK-state machine
Decode Ways (LC 91)LeetCode2-state machine
Domino and Tromino Tiling (LC 790)LeetCodeProfile state machine

State machine DP transforms confusing problems into clear diagrams. Draw the states, draw the arrows, and the code writes itself. Every stock problem — and many beyond — follows this pattern.