State Design

For many advanced DP problems, simply using an index i or (i, j) is not enough. You often need to track additional information, like “Am I currently holding a stock?” or “Is the previous element included?”.

The best way to model these problems is using a State Machine.

1. Case Study: Best Time to Buy and Sell Stock with Cooldown

Problem Statement: You want to maximize your profit by buying and selling a stock.

  • You can complete as many transactions as you like.
  • After you sell your stock, you cannot buy stock on the next day (cooldown of 1 day).

Defining the States

We can define three states for any given day:

  1. Held: We are holding a stock.
  2. Sold: We just sold a stock (we are in the cooldown period).
  3. Reset: We are not holding a stock and not in cooldown (ready to buy).

Transitions

  • From Held:
  • Sell → Go to Sold (Profit increases)
  • Wait → Stay in Held
  • From Sold:
  • Wait → Go to Reset (Cooldown over)
  • From Reset:
  • Buy → Go to Held (Profit decreases)
  • Wait → Stay in Reset

Interactive State Machine

Simulate the trading process. Try to maximize your profit given the price history.

Prices: Day 1: 1, Day 2: 2, Day 3: 3, Day 4: 0, Day 5: $2

[!TIP] Try it yourself: Click “Buy”, “Sell”, or “Wait” to navigate the state machine. Watch how your profit changes.

Current Day: 1 / 5
Profit: $0
Current Price: $1
Reset
Held
Sold
Log: Game started.

Recurrence Relation

  • held[i]: Max profit on day i ending in Held state.
  • Either kept holding: held[i-1]
  • Or bought today: reset[i-1] - price[i]

  • held[i] = max(held[i-1], reset[i-1] - price[i])

  • sold[i]: Max profit on day i ending in Sold state.
  • Must have sold today: held[i-1] + price[i]

  • sold[i] = held[i-1] + price[i]

  • reset[i]: Max profit on day i ending in Reset state.
  • Either kept resting: reset[i-1]
  • Or cooled down from sold: sold[i-1]
  • reset[i] = max(reset[i-1], sold[i-1])

Code Implementation

// Java
public class StockCooldown {
  public int maxProfit(int[] prices) {
    if (prices.length <= 1) return 0;

    int held = -prices[0];
    int sold = 0;
    int reset = 0;

    for (int i = 1; i < prices.length; i++) {
      int prevHeld = held;
      int prevSold = sold;
      int prevReset = reset;

      // Update states based on transitions
      held = Math.max(prevHeld, prevReset - prices[i]);
      sold = prevHeld + prices[i];
      reset = Math.max(prevReset, prevSold);
    }

    return Math.max(sold, reset);
  }
}

// Go
func MaxProfit(prices []int) int {
  if len(prices) <= 1 {
    return 0
  }

  held := -prices[0]
  sold := 0
  reset := 0

  for i := 1; i < len(prices); i++ {
    prevHeld := held
    prevSold := sold
    prevReset := reset

    held = max(prevHeld, prevReset - prices[i])
    sold = prevHeld + prices[i]
    reset = max(prevReset, prevSold)
  }

  return max(sold, reset)
}

func max(a, b int) int {
  if a > b {
    return a
  }
  return b
}

2. Deep Dive Strategy Lab: State Design

Intuition Through Analogy

Think of this chapter like budget planning across months. The goal is not to memorize a fixed trick, but to repeatedly answer:

  1. What is the bottleneck?
  2. Which constraint dominates (time, memory, latency, correctness)?
  3. Which representation makes the bottleneck easier to eliminate?