Stock trading with k transactions and cooldown in Python

Why this matters

Many real problems ask you to make a choice each day, where today's choice depends on what you did before. A stock trader either holds a share or does not, and whether the trader can buy tomorrow may depend on whether a sale happened today. Dynamic programming (DP, a method that solves a problem by reusing answers to smaller versions of it) handles this well when you describe each day by a small set of states.

A state is a label for the situation you are in at the end of a day. This lesson builds the states for a stock problem with two rules: a limit on the number of trades, and a cooldown after each sale. The previous lesson, burst balloons, uses a different shape of DP, where the answer depends on ranges of positions. Here the answer depends on the day and the state.

The problem

You are given a list of prices, one for each day. A list is Python's ordered collection of values, written with square brackets, such as [1, 2, 3, 0, 2]. Each day you may do one of three things: buy one share, sell the share you hold, or do nothing. You can hold at most one share at a time. The goal is to finish with no share in hand and as much money as possible.

Two extra rules make it harder.

• A trade limit. You may make at most k transactions. A transaction is one buy followed later by one sell.
• A cooldown. After you sell on a day, you cannot buy on the next day.

Take the prices [1, 2, 3, 0, 2] with a cooldown and no limit on trades. Here is a plan that earns 3.

1. Day 0, price 1: buy.
2. Day 1, price 2: sell. Profit so far is 1.
3. Day 2, price 3: cooldown, because you sold yesterday. Do nothing.
4. Day 3, price 0: buy. Allowed, because the day before was a cooldown day, not a sale day.
5. Day 4, price 2: sell. Profit is 1 + 2 = 3.

A plan that buys on day 2 right after the sale on day 1 breaks the cooldown rule. The DP has to find the best plan without trying every one.

Why one number per day is not enough

A first idea is to store one number per day: the best money you can have on that day. This does not work, because two plans can leave you with the same money and different options for tomorrow. One plan may hold a share and another may have no share. The first plan can sell tomorrow. The second can buy tomorrow. Those options differ, so the DP needs to track them separately.

The fix is to keep one number per state. The states for this problem are:

• hold: you hold one share at the end of the day. The number is your money, which is smaller because you paid for the share.
• sold: you sold a share on this day. The number is your money. Tomorrow you cannot buy because of the cooldown.
• rest: you hold no share, and you did not sell today. Tomorrow you may buy.

Each state stores the best money you can have in that state at the end of the day. The word "best" is the important part: for each state, keep the largest value and throw away the others.

The three transitions

A transition says how today's states depend on yesterday's states and today's price p.

To be holding at the end of today, you either held yesterday and did nothing, or you did not hold yesterday and bought today. You cannot buy today if you sold yesterday, so the buy comes from rest yesterday. The rule is:

hold_today = max(hold_yesterday, rest_yesterday - p)

To have sold today, you must have held yesterday and sold today at price p. The rule is:

sold_today = hold_yesterday + p

To be resting at the end of today, you did not hold a share and you did not sell today. You either were resting yesterday, or you sold yesterday, which means today is the cooldown day. The rule is:

rest_today = max(rest_yesterday, sold_yesterday)

The starting point is the day before the first price. You have no share and no sale, so rest is 0. The states hold and sold are not possible yet, so they start at a value far below any real amount of money. The final answer is the best of sold and rest on the last day. hold is not a valid ending, because it means a share is still unsold.

Worked example: the prices 1, 2, 3, 0, 2

Here is the table of the three values after each day. A dash means the state is not possible yet. The numbers came from a Python program that applies the three transitions.

start   hold=- sold=- rest=0
day 0 p=1 hold=-1 sold=- rest=0
day 1 p=2 hold=-1 sold=1 rest=0
day 2 p=3 hold=-1 sold=2 rest=1
day 3 p=0 hold=1 sold=-1 rest=2
day 4 p=2 hold=1 sold=3 rest=2

Read day 1. sold is 1 because you held a share from day 0 at price 1, and sold it at price 2. Then rest on day 2 is 1, because rest takes the larger of yesterday's rest (0) and yesterday's sold (1).

Read day 3. hold becomes 1. The rule says max(hold_yesterday, rest_yesterday - p). Yesterday's hold was -1, and yesterday's rest was 1, so rest - p is 1 - 0 = 1. The bigger value is 1. This is the buy at price 0 on day 3, made from a rest value that comes from a sale on day 1 with a cooldown on day 2.

Read the last day. sold is 3, which is hold from day 3 (value 1) plus the price 2. The best finish is max(sold, rest), which is max(3, 2), so the answer is 3.

Notice that rest on day 3 is 2, not 1. It takes the sold value of 2 from day 2, which is the sale at price 3. That is allowed, because rest on day 3 is used for a buy on day 4, and the cooldown only blocks the day right after a sale.

Using three variables

The table only ever needs yesterday's values to compute today's. So the code can keep three variables and overwrite them each day. The memory does not grow with the number of days.

Save the functions in this lesson in one file named profit.py.

NEG = float("-inf")
 
 
def max_profit_cooldown(prices):
    hold = NEG
    sold = NEG
    rest = 0
    for p in prices:
        hold, sold, rest = max(hold, rest - p), hold + p, max(rest, sold)
    return max(sold, rest)

NEG is a name for negative infinity. float("-inf") makes a floating-point number (a number with a decimal point) that is smaller than every real number. Adding or subtracting a price keeps it at negative infinity, and max never picks it over a real amount of money. If you know Java, this is the place where Java code needs Integer.MIN_VALUE / 2 to avoid wrap-around. Python integers do not wrap around, so this problem does not arise.

The line with three names on the left is a tuple assignment. A tuple is a group of values that Python handles together, here written without brackets. Python first works out all three values on the right, using yesterday's values. Only then does it store them. That gives the same result as computing new values first. A version that stores one value before computing the next is shown in the mistakes section.

The function returns max(sold, rest). With the prices [1, 2, 3, 0, 2], it returns 3.

Worked example 2: at most k transactions

Now remove the cooldown and keep the limit. The number of transactions is limited to k, and each transaction is one buy and one later sell.

Take [2, 4, 1] with k equal to 2. The best plan is to buy at 2 and sell at 4, which is a profit of 2. Then buying at 1 gives nothing more, because there is no later price to sell at. The answer is 2.

Take [3, 2, 6, 5, 0, 3] with k equal to 2. Buy at 2 and sell at 6 for a profit of 4. Then buy at 0 and sell at 3 for a profit of 3. The total is 7. The two transactions are separate, and you cannot hold two shares at once, so this is the most you can do.

The three-variable trick does not work here, because the number of transactions matters. The DP needs one value for each number of transactions finished so far. Two lists do the job:

• buy[j]: the best money you can have while holding a share, after starting the jth transaction with a buy.
• sell[j]: the best money you can have after finishing j transactions with a sell.

For each day with price p, and for each j from 1 to k:

buy[j]  = max(buy[j], sell[j - 1] - p)
sell[j] = max(sell[j], buy[j] + p)

The first line says: to start transaction j, you need the money from the end of transaction j - 1, and you pay p now. The second line says: to finish transaction j, you must have bought it, and you receive p now.

Here is the code. It goes in profit.py below the functions above.

def max_profit_k(k, prices):
    buy = [NEG] * (k + 1)
    sell = [0] * (k + 1)
    for p in prices:
        for j in range(1, k + 1):
            buy[j] = max(buy[j], sell[j - 1] - p)
            sell[j] = max(sell[j], buy[j] + p)
    return max(sell)

[NEG] * (k + 1) makes a list of k + 1 copies of NEG. The sell list starts at 0, which means no money earned yet. max(sell) is the best finish over every transaction count, and sell[0] is always 0, so it covers the case of no trades.

The memory is two lists of size k + 1, which is 2k + 2 numbers. It does not depend on the number of days. A table with one row per day and one column per transaction would use n times k numbers, so the lists save a lot when the list of prices is long.

The inner loop goes from 1 up to k. When it reaches j, the value sell[j - 1] has already been updated for today if j - 1 is at least 1. That lets a sale on today's price feed the buy for the next transaction on the same day. The brute force check described below, which tries every option, agrees with this code on 1,500 random lists.

Worked example 3: cooldown and a limit together

Combine the two rules. Each day now needs hold, sold and rest for every transaction count. Each of the three is a list of size k + 1. The index j is the number of buys made so far.

The transitions are the same as in the first section, with one change. A buy moves the count from j - 1 to j, so the buy uses rest[j - 1] and writes hold[j].

hold_today[j] = max(hold[j], rest[j - 1] - p)
sold_today[j] = hold[j] + p
rest_today[j] = max(rest[j], sold[j])

Each day's new values are computed into fresh lists, and then the names are moved to them. Writing into the old lists would mix yesterday's values with today's, the same bug as in the three-variable code.

def max_profit_k_cooldown(k, prices):
    hold = [NEG] * (k + 1)
    sold = [NEG] * (k + 1)
    rest = [0] * (k + 1)
    for p in prices:
        new_hold = [NEG] * (k + 1)
        new_sold = [NEG] * (k + 1)
        new_rest = [0] * (k + 1)
        for j in range(k + 1):
            if j >= 1:
                new_hold[j] = max(hold[j], rest[j - 1] - p)
                new_sold[j] = hold[j] + p
            new_rest[j] = max(rest[j], sold[j])
        hold, sold, rest = new_hold, new_sold, new_rest
    return max(max(sold), max(rest))

The rest list starts at 0 for every count. That is safe because 0 means doing nothing, and a limit of j buys allows fewer.

Checking the three methods

A brute force check is a second program that tries every sequence of actions on small price lists, one action per day. It is slow, but it follows the rules directly, so it is a good test. The function below tries every choice on each day. It has an inner function named go, which is a function defined inside another function. go can read prices, k and cooldown from the outer function.

def brute(prices, k, cooldown):
    def go(i, holding, buys, sold_yesterday, cash):
        if i == len(prices):
            return NEG if holding else cash
        p = prices[i]
        best = go(i + 1, holding, buys, False, cash)
        if holding:
            best = max(best, go(i + 1, False, buys, True, cash + p))
        elif buys < k and not (cooldown and sold_yesterday):
            best = max(best, go(i + 1, True, buys + 1, False, cash - p))
        return best
 
    return go(0, False, 0, False, 0)

The brute force returns the money at the end, but only when no share is held. Its parameter sold_yesterday is the cooldown rule written out in code.

The test code below calls the functions from one main function. Put it at the bottom of profit.py and run the file with python3 profit.py:

def main():
    print(max_profit_cooldown([1, 2, 3, 0, 2]))
    print(max_profit_k(2, [2, 4, 1]))
    print(max_profit_k(2, [3, 2, 6, 5, 0, 3]))
    print(max_profit_k_cooldown(2, [3, 2, 6, 5, 0, 3]))
    print(max_profit_k_cooldown(1, [1, 2, 3, 0, 2]))
    rnd = random.Random(7)
    checked = 0
    for _ in range(1500):
        prices = [rnd.randrange(10) for _ in range(rnd.randrange(9))]
        k = rnd.randrange(4)
        ok = (max_profit_cooldown(prices) == brute(prices, len(prices) + 1, True)
              and max_profit_k(k, prices) == brute(prices, k, False)
              and max_profit_k_cooldown(k, prices) == brute(prices, k, True))
        if not ok:
            print("MISMATCH", prices)
            return
        checked += 1
    print("random checks passed:", checked)
 
 
if __name__ == "__main__":
    main()

profit.py also needs import random at the top, which the test uses to make random lists. The brute call for the cooldown-only method passes len(prices) + 1 as the limit, which is more than enough transactions. The and between the three checks means the program stops at the first one that is false. rnd.randrange(9) gives a length from 0 to 8, so the list can be empty.

Its output is:

3
2
7
7
2
random checks passed: 1500

These lines come from five calls. The first is the cooldown-only method on [1, 2, 3, 0, 2], which gives 3. The second is the limited method on [2, 4, 1] with k equal to 2, which gives 2. The third is the limited method on [3, 2, 6, 5, 0, 3], which gives 7. The fourth is the combined method with k equal to 2 on the same prices, which also gives 7. The fifth is the combined method with k equal to 1 on [1, 2, 3, 0, 2], which gives 2.

Common mistakes

Forgetting the cooldown

If the buy line uses hold or rest from the wrong day, the cooldown is lost. Using rest - p for the buy is correct only because rest does not include a sale from yesterday. If the buy line uses max(rest, sold) - p instead, a sale is followed by a buy on the very next day, which the cooldown forbids.

Here is the cooldown-only function with that change:

NEG = float("-inf")
 
 
def cooldown_buy_after_sale(prices):
    hold, sold, rest = NEG, NEG, 0
    for p in prices:
        hold, sold, rest = max(hold, max(rest, sold) - p), hold + p, max(rest, sold)
    return max(sold, rest)
 
 
print(cooldown_buy_after_sale([1, 2, 3, 0, 2]))

Its output is 4, where the correct answer is 3. The plan it finds buys at 1 on day 0, sells at 3 on day 2, buys at 0 on day 3, and sells at 2 on day 4. The buy on day 3 comes right after the sale on day 2, which the rules do not allow.

Nothing crashes, and the output is a plausible number, so this bug only shows up when you compare against a case you know well.

Starting holding and sold at zero

The starting value for hold and sold must be very small, so that the DP cannot use a share it never paid for. A common mistake is to write zero:

def cooldown_zero_start(prices):
    hold, sold, rest = 0, 0, 0
    for p in prices:
        hold, sold, rest = max(hold, rest - p), hold + p, max(rest, sold)
    return max(sold, rest)
 
 
print(cooldown_zero_start([1, 2, 3, 0, 2]))

Its output is 4, but the correct answer is 3. With a start of zero, hold looks like you already own a share at no cost. The extra profit comes from a share that was never bought.

Using None as the start value

Python has a value named None, which means "no value". It looks like a good start for hold and sold, but it cannot be used in arithmetic or comparisons:

def cooldown_none_start(prices):
    hold, sold, rest = None, None, 0
    for p in prices:
        hold, sold, rest = max(hold, rest - p), hold + p, max(rest, sold)
    return max(sold, rest)
 
 
print(cooldown_none_start([1, 2, 3, 0, 2]))

Python stops with this error:

Traceback (most recent call last):
  File "/path/to/none_start.py", line 8, in <module>
    print(cooldown_none_start([1, 2, 3, 0, 2]))
          ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/path/to/none_start.py", line 4, in cooldown_none_start
    hold, sold, rest = max(hold, rest - p), hold + p, max(rest, sold)
                       ~~~^^^^^^^^^^^^^^^^
TypeError: '>' not supported between instances of 'int' and 'NoneType'

The path depends on where you saved the file. The message says max tried to compare a number with None, which has no order. The fix is float("-inf"), as in the code above.

Changing values in place

If the three lines are written one after another, each line sees the values that the lines above it already changed:

NEG = float("-inf")
 
 
def cooldown_sequential(prices):
    hold, sold, rest = NEG, NEG, 0
    for p in prices:
        hold = max(hold, rest - p)
        sold = hold + p
        rest = max(rest, sold)
    return max(sold, rest)
 
 
print(cooldown_sequential([1, 2, 3, 0, 2]))

Its output is 4. The line for rest now uses today's sold, not yesterday's. So a sale counts as a rest day on the same day, and the buy on day 3 breaks the cooldown. The tuple assignment in the correct version avoids this, because all three values are worked out from yesterday's values first.

Returning the wrong state at the end

The answer must come from sold or rest, never from hold, because hold means a share is still in hand. Returning hold would count a share that was never sold. The code uses max(sold, rest), and the combined version takes the maximum over every transaction count.

Practice

Exercise 1

Find the answer for [6, 1, 3, 2, 4, 7] with a cooldown and no limit on trades. Explain which days you buy and sell.

Show the answer

The answer is 6. Buy at price 1 on day 1, and sell at price 7 on day 5. The profit is 6.

Another plan is to buy at 1 on day 1, sell at 3 on day 2, skip day 3 because of the cooldown, buy at 4 on day 4, and sell at 7 on day 5. That earns 2 + 3 = 5, which is less than 6.

Exercise 2

For the prices [1, 2, 3] with a cooldown, what are hold, sold and rest after the last day?

Show the answer

Use the transitions. After day 0 (price 1), hold is -1, sold is not possible, and rest is 0. After day 1 (price 2), hold is -1, sold is 1, and rest is 0. After day 2 (price 3), hold is -1, sold is 2 (the share from day 0 sold at 3), and rest is 1. The final answer is max(2, 1), which is 2.

Exercise 3

With a limit of 5 transactions, how many numbers does the list method keep while it runs? Count only the numbers that change during the work.

Show the answer

With k equal to 5 the lists have size 6. The entries buy[1] through buy[5] and sell[1] through sell[5] change during the work. That is 10 numbers. The entry sell[0] stays 0, and buy[0] is never used. The memory does not grow with the number of days.

Exercise 4

Run max_profit_k with k equal to 5 on [3, 2, 6, 5, 0, 3]. Why is the answer still 7?

Show the answer

The function returns 7 because the prices allow only two useful transactions. The buy at 2 and sell at 6 gives 4, and the buy at 0 and sell at 3 gives 3. After the sale at 3 on day 5, there are no more days left, so no third transaction fits. Raising the limit cannot add profit that the prices do not offer.

Exercise 5

Explain in plain words why the cooldown-only method needs three values and not two.

Show the answer

Two values are not enough, because two plans can leave the same money with different options. One plan may hold a share, and another may have just sold one. A day after a sale, the cooldown blocks a buy, while a day after a rest it does not. The state sold records that a sale happened today, and rest records that the trader is free to buy tomorrow. hold records that a share is in hand. Each needs its own value.

Summary

• Describe each day by states. For a cooldown, the states are holding, just sold, and resting.
• The transitions use yesterday's values only, so three variables are enough for the cooldown.
• To limit transactions, keep one buy and one sell value for each count from 1 to k.
• For both rules together, keep lists of size k + 1 for each state, and build new lists each day.
• Start the holding and sold states at negative infinity, and never let the answer come from the holding state.
• Check against a brute force that tries every sequence of actions on small inputs.