Stock trading with k transactions and cooldown

Why this matters

A common question in programming is how to choose actions over time when some actions block others. A game character can carry one item at a time. A user can make a purchase only after the previous payment clears. A machine can run one job, then must rest before the next. Each choice depends on what was done before, and the best answer needs the whole history.

Dynamic programming handles these cases by keeping a small set of states. A state is a label for the situation you are in right now. If the labels are chosen well, the best answer for each label can be computed from the labels of the day before. This lesson builds that method on a stock trading problem. The idea of storing answers in a table comes from the memoization and tabulation lessons in this module. The problem here has a table over days, and the lesson shows how to shrink it.

You need C++ std::vector, long long, std::max and loops from earlier modules. The problem uses only simple arithmetic, so no other lesson in the module is needed.

The problem in plain words

You are given a list of prices, one per day. Each price is the price of one share on that day. You may hold at most one share at a time. A transaction is one buy followed by one sell. The goal is to finish with no shares and the largest total cash.

There are two rules on top of that.

• A limit k says you may make at most k transactions. The limit is part of some versions of the problem and not others. When there is no limit, you can make as many as you like.
• A cooldown says that after you sell a share on day d, you cannot buy on day d + 1. Day d + 1 is the cooldown day. You can buy again on day d + 2.

Take the prices 1, 2, 3, 0, 2 for five days, numbered 0 to 4, with a cooldown and no limit. A good plan is this.

• Day 0: buy at price 1.
• Day 1: sell at price 2. Profit so far is 1.
• Day 2: cooldown. You cannot buy today, so you do nothing.
• Day 3: buy at price 0.
• Day 4: sell at price 2. Profit is 1 + 2 = 3.

The total is 3. Without the cooldown, the best plan would be to buy at 1, sell at 3 on day 2, then buy at 0 on day 3 and sell at 2 on day 4. That plan makes 2 + 2 = 4. The cooldown rule forbids buying on day 3 after the sale on day 2, so the answer with cooldown is 3.

Two things to keep in mind. The answer is cash, not shares, so any share still held at the end does not count. And a day can have only one action: buy, sell or nothing.

Brute force first

Before any table, think about the direct way. On each day, choose one of three actions: do nothing, buy, or sell. Keep track of whether you hold a share, whether you sold yesterday, and how many buys you have made. At the end, add up the cash. Keep the best total among all the action lists that follow the rules and end with no share.

With n days and three actions each day, there are 3 * 3 * ... * 3, which is 3 to the power n, action lists to check. For 5 days that is 243. For 20 days that is about 3.5 billion. This is too slow for long price lists, but it is a good check on small ones, and the program below uses it that way.

The direct approach is correct but slow for the same reason as in other DP lessons: the same situation appears again and again. Many different action lists end up on day 5 holding a share with the same cash. Dynamic programming stores the best cash for each situation once, so those repeats cost nothing.

The idea: three states on each day

At the end of every day, you are in exactly one of three states. For each state, store the best cash you can have in that state.

• hold: you own one share. The stored number is the cash you have, which is usually negative, because you paid for the share.
• sold: you sold a share today. The stored number is the cash you have. Tomorrow is a cooldown day, so you cannot buy tomorrow.
• rest: you own no share, and you are free to buy tomorrow. The stored number is the cash you have.

These three states cover every situation. A state machine is a set of states with rules that say how you move from one to another. The table below shows the moves for one day with price p.

The table leaves out one more move, the cooldown. The state sold does not allow a buy. On the next day, a sold state moves to rest, and only rest can buy. The rule is written this way: the buy move starts from rest, so a day that starts in sold cannot buy. This is the entire cooldown rule, and it costs nothing extra in the code.

Each day you compute the three new values from the three old values, using only the moves in the table. The new value for a state is the best of all the moves that end in that state.

• New hold: the better of keeping the old hold, or buying today from yesterday's rest. Buying gives rest - p.
• New sold: the old hold, plus today's price p. Only a hold can become a sale.
• New rest: the better of keeping yesterday's rest, or moving yesterday's sold into rest. A sold from yesterday becomes rest today, which is the cooldown ending.

The three formulas are the whole method. The rest of the lesson shows how to run them.

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

Start on day 0 before any action. You hold nothing and have no cash, so the rest state is 0. The hold and sold states are impossible at that point, so their stored value is minus infinity. In the code, minus infinity is a very large negative number, so that impossible states never win a comparison.

The table below shows the three stored values at the end of each day. The first column gives the day and its price. The values are the output of the program in the next section.

Read the rows one by one.

• Day 0: the buy at 1 gives hold = 0 - 1 = -1. Nothing else is possible yet, so sold stays impossible and rest stays 0.
• Day 1: the sale at 2 gives sold = -1 + 2 = 1. The hold value stays -1, because keeping it is better than buying from rest for 2, which would give 0 - 2 = -2.
• Day 2: the sale at 3 gives sold = -1 + 3 = 2. Rest becomes 1, because the sold value from day 1, which is 1, moves into rest today.
• Day 3: buying at price 0 from rest gives hold = 1 - 0 = 1. The rest value 1 came from the sale on day 1, and day 2 was the cooldown. Day 3 is the first day that can buy. Selling a share at 0 gives sold = -1. The rest value becomes 2, because the sold value 2 from day 2 moves into rest.
• Day 4: the sale at 2 gives sold = 1 + 2 = 3. Rest stays 2, which is better than the sold value -1 from day 3.

The answer is the best of the two states that hold no share at the end, rest and sold. At day 4 that is max(2, 3), which is 3. The hold state is not counted, because it holds a share.

The code with three variables

Because each day only needs yesterday's three values, the program does not need a table over days. It keeps three variables and updates them once per day. This is the saving that the lesson summary promises, so the memory is three numbers no matter how many days there are.

The update has one subtle point. The three new values must all use yesterday's values. If the code writes hold before it reads hold again for sold, the wrong number is used. The program stores the three new values in temporary variables first, then copies them over.

The program below also has a brute force function that uses only the rules, and random tests that compare the two. Every case uses the same price lists and checks that the answers match.

#include <algorithm>
#include <iostream>
#include <limits>
#include <random>
#include <vector>
 
// Marks an impossible state. Large enough to never win a max, small enough
// that subtracting prices from it cannot overflow.
constexpr long long NEG = -(1LL << 60);
 
// At most k buys, one-day cooldown after each sale.
// hold[b]: holding a share, b buys made so far.
// sold[b]: sold today, b buys made so far (tomorrow is cooldown).
// rest[b]: not holding and free to buy tomorrow, b buys made so far.
// The loop goes from high b to low b, so rest[b - 1] still has the value
// from the previous day when it is read.
long long maxProfitK(int k, const std::vector<long long>& prices) {
    if (k <= 0 || prices.empty()) return 0;
    std::vector<long long> hold(k + 1, NEG), sold(k + 1, NEG), rest(k + 1, NEG);
    rest[0] = 0;
    for (long long p : prices) {
        for (int b = k; b >= 0; --b) {
            const long long oldHold = hold[b];
            const long long oldSold = sold[b];
            const long long oldRest = rest[b];
            long long newHold = oldHold;
            if (b >= 1) newHold = std::max(newHold, rest[b - 1] - p);
            hold[b] = newHold;
            sold[b] = oldHold + p;
            rest[b] = std::max(oldRest, oldSold);
        }
    }
    return std::max(*std::max_element(rest.begin(), rest.end()),
                    *std::max_element(sold.begin(), sold.end()));
}
 
// No limit on transactions: the b dimension disappears.
long long maxProfitUnlimited(const std::vector<long long>& prices) {
    long long hold = NEG, sold = NEG, rest = 0;
    for (long long p : prices) {
        const long long newHold = std::max(hold, rest - p);
        const long long newSold = hold + p;
        const long long newRest = std::max(rest, sold);
        hold = newHold;
        sold = newSold;
        rest = newRest;
    }
    return std::max(rest, sold);
}
 
// Brute force from the rules only: each day is nothing, buy or sell.
long long bruteForce(int k, const std::vector<long long>& prices) {
    const int n = static_cast<int>(prices.size());
    long long best = 0;
    long long total = 1;
    for (int i = 0; i < n; ++i) total *= 3;
    for (long long code = 0; code < total; ++code) {
        long long c = code;
        bool holding = false;
        bool soldYesterday = false;
        int buys = 0;
        long long cash = 0;
        bool valid = true;
        for (int day = 0; day < n && valid; ++day) {
            const int action = static_cast<int>(c % 3);
            c /= 3;
            bool soldToday = false;
            if (action == 1) {                         // buy
                if (holding || soldYesterday || buys >= k) valid = false;
                else { holding = true; ++buys; cash -= prices[day]; }
            } else if (action == 2) {                  // sell
                if (!holding) valid = false;
                else { holding = false; cash += prices[day]; soldToday = true; }
            }
            soldYesterday = soldToday;
        }
        if (valid && !holding) best = std::max(best, cash);
    }
    return best;
}
 
int main() {
    std::cout << "k=2 [2,4,1] -> " << maxProfitK(2, {2, 4, 1}) << "\n";
    std::cout << "k=2 [3,2,6,5,0,3] -> " << maxProfitK(2, {3, 2, 6, 5, 0, 3}) << "\n";
    std::cout << "unlimited [1,2,3,0,2] -> " << maxProfitUnlimited({1, 2, 3, 0, 2}) << "\n";
    std::cout << "unlimited [1] -> " << maxProfitUnlimited({1}) << "\n";
    std::cout << "k=1 [1,2,3,0,2] -> " << maxProfitK(1, {1, 2, 3, 0, 2}) << "\n";
    std::cout << "k=3 [1,2,3,0,2] -> " << maxProfitK(3, {1, 2, 3, 0, 2}) << "\n";
 
    std::mt19937 rng(12345);
    std::uniform_int_distribution<int> lenDist(0, 7), priceDist(0, 9), kDist(1, 3);
    int checked = 0, failures = 0;
    for (int t = 0; t < 3000; ++t) {
        const int n = lenDist(rng);
        std::vector<long long> prices(n);
        for (auto& p : prices) p = priceDist(rng);
        const int k = kDist(rng);
        const long long want = bruteForce(k, prices);
        const long long got = maxProfitK(k, prices);
        ++checked;
        if (want != got) ++failures;
        // With k larger than the number of days, the limit never applies.
        if (maxProfitUnlimited(prices) != maxProfitK(100, prices)) ++failures;
    }
    std::cout << "random checks=" << checked << " failures=" << failures << "\n";
}

Compile and run it with:

g++ -std=c++20 -Wall -o stocks stocks.cpp
./stocks

The program prints the following. The compile produced no warnings. The programs in this lesson were compiled with g++ 15.2.0.

k=2 [2,4,1] -> 2
k=2 [3,2,6,5,0,3] -> 7
unlimited [1,2,3,0,2] -> 3
unlimited [1] -> 0
k=1 [1,2,3,0,2] -> 2
k=3 [1,2,3,0,2] -> 3
random checks=3000 failures=0

The first line is a short list with a limit of 2, which gives 2 by buying at 2 and selling at 4. The second line is a longer list with a limit of 2, which gives 7 by buying at 2 and selling at 6, then buying at 0 and selling at 3. The third and fourth lines are the unlimited version. The fifth and sixth lines show that a limit of 1 gives 2, while a limit of 3 gives 3, the same as no limit, because the prices 1, 2, 3, 0, 2 never need more than two transactions.

The random check ran 3,000 price lists with 0 to 7 days. Each list was checked against the brute force for a random limit from 1 to 3, and against the unlimited version. All 3,000 agreed.

Limiting the transactions

With a limit k, the state needs one more number: how many buys you have made so far. The rule is that a buy is allowed only when the count is below k. The state becomes a count for each of the three states. Each state then has one stored value per count, so the program keeps three arrays of length k + 1.

The loop still runs over the days, and inside each day it runs over the count b. The formulas are the same as before, with one change. A buy moves from a rest state with b - 1 buys to a hold state with b buys, so the buy uses rest[b - 1]. A sale does not change the count.

The loop over b runs from high to low. This is the same ordering point as in the three-variable code. When the program updates hold[b], it must still read the old rest[b - 1]. Going down from k means rest[b - 1] has not been changed yet on that day. The brute force tests checked that this ordering gives the right answers.

The table over days is gone. The memory is three arrays of k + 1 numbers, so it depends on k and not on the number of days. A table over days and counts would store n times k values for each state, which grows with the length of the price list.

Worked example 2 uses the limit to show the difference. With k = 2 and the prices 3, 2, 6, 5, 0, 3, the program prints 7. The best plan is to buy at 2 on day 1 and sell at 6 on day 2, which makes 4. Then buy at 0 on day 4 and sell at 3 on day 5, which makes 3. The total is 4 + 3, which is 7. Day 3 is the cooldown day after the sale on day 2, so the buy on day 4 is two days after the sale, which is allowed.

Worked example 3: the same prices with k = 1

With k = 1, only one transaction is allowed. For the prices 1, 2, 3, 0, 2, the best single plan is to buy at 1 on day 0 and sell at 3 on day 2, which makes 2. The program prints 2 for that case, as shown in the output above.

This example shows what the limit does. With no limit, the best plan is to buy at 1 on day 0 and sell at 2 on day 1, which makes 1, then buy at 0 on day 3 and sell at 2 on day 4, which makes 2. The total is 3. The cooldown on day 2 is what allows the buy on day 3. The limit of 1 removes the second transaction, and the single transaction that remains is worth 2. A reader can check each of these by hand, which is why small examples like this are useful before trying a long list.

Common mistakes

Forgetting the cooldown. If the state machine has only hold and rest, the buy can happen on the day after a sale. For the prices 1, 2, 3, 0, 2, the answer without cooldown is 4, because it buys at 1, sells at 3 on day 2, and buys at 0 on day 3. The cooldown forbids that buy, so the correct answer is 3. The sold state is there to block the buy.

Counting a share that is still held. The answer must be a state with no share, which is rest or sold. A common shortcut is to take the best value over all three states, including hold. The hold value is the cash after paying for a share that has not been sold, so it does not describe a finished plan. The code returns the best of rest and sold only.

Updating the states in the wrong order. When the arrays are updated in place, a loop that goes upward reads rest[b - 1] after it has already been changed on the same day. The next program repeats the limit version with the loop going upward. It runs without an error but gives a different answer.

// Same as maxProfitK in the lesson, except that the loop over b goes upward.
long long wrongOrderK(int k, const std::vector<long long>& prices) {
    std::vector<long long> hold(k + 1, NEG), sold(k + 1, NEG), rest(k + 1, NEG);
    rest[0] = 0;
    for (long long p : prices) {
        for (int b = 0; b <= k; ++b) {
            const long long oldHold = hold[b];
            const long long oldSold = sold[b];
            const long long oldRest = rest[b];
            long long newHold = oldHold;
            if (b >= 1) newHold = std::max(newHold, rest[b - 1] - p);  // rest[b - 1] may be new already
            hold[b] = newHold;
            sold[b] = oldHold + p;
            rest[b] = std::max(oldRest, oldSold);
        }
    }
    return std::max(*std::max_element(rest.begin(), rest.end()),
                    *std::max_element(sold.begin(), sold.end()));
}

Calling it with k = 3 on the prices 1, 2, 3, 0, 2 prints 4, while the correct answer is 3. The loop went upward, so on day 3 the buy at 0 read a rest[b - 1] that already included that same day's update.

Using int for the impossible value. If the impossible state is stored as the smallest int, one subtraction of a price overflows. The next program does that with int, and the compiler with -fsanitize=undefined reports the overflow.

#include <climits>
#include <iostream>
 
int main() {
    // A plain int sentinel for "impossible" and one subtraction from it.
    int hold = INT_MIN;
    int price = 1;
    std::cout << (hold - price) << "\n";
}

Compiled with g++ -std=c++20 -Wall -fsanitize=undefined -o overflow overflow.cpp, the run prints this.

overflow.cpp:8:36: runtime error: signed integer overflow: -2147483648 - 1 cannot be represented in type 'int'
2147483647

The second line is the wrapped value. Signed overflow is undefined behaviour in C++, which means the program may do anything, so the result cannot be trusted. The fix is the one the lesson uses: store the prices and the states as long long, and pick an impossible value that is far from any real cash amount, not the smallest possible number.

Practice

Exercise 1. Find the best cash with no limit and with cooldown for the prices 7, 1, 5, 3, 6, 4. Name the buy and sell days.

Show the answer. The best plan is to buy at 1 on day 1 and sell at 6 on day 4, which makes 5. Buying at 1 and selling at 5 on day 2 makes 4, but then day 3 is the cooldown, and the next buy would be on day 4 at price 6, which cannot be sold later at a profit. So the best total is 5. The program maxProfitUnlimited returns 5 for this list.

Exercise 2. With the limit k = 1, find the best cash for the prices 3, 3, 5, 0, 0, 3, 1, 4.

Show the answer. Buy at 0 on day 3 or day 4, and sell at 4 on day 7. The profit is 4. A limit of one transaction means no other buy and sell pair can be added, so the answer is 4. The function maxProfitK returns 4 for this list when k is 1.

Exercise 3. For the prices 1, 2, with no limit and cooldown, write the three state values at the end of day 0 and at the end of day 1, then give the answer.

Show the answer. At the end of day 0, hold is -1, sold is impossible, and rest is 0. At the end of day 1, the sale at 2 gives sold = -1 + 2 = 1. Hold stays -1, and rest stays 0, because the old sold value is impossible. The answer is max(rest, sold) = max(0, 1) = 1.

Exercise 4. With k = 2, find the best cash for the prices 1, 2, 3, 4, 5.

Show the answer. One transaction, buy at 1 and sell at 5, makes 4. Two transactions cannot make more here, because the prices only go up. One buy at the start and one sell at the end already collects every rise, and splitting that into two transactions adds nothing. The function maxProfitK returns 4 for this list when k is 2.

Exercise 5. After a sale on day 4, which is the first day you can buy again? Then say how many action lists the brute force checks for 8 days.

Show the answer. Day 5 is the cooldown, so the first buy is on day 6. For 8 days, each day has three actions, so the brute force checks 3 to the power 8, which is 6,561 lists. This is why the brute force is only used for short lists.

Summary

• Give each day three states: hold a share, sold today, or rest with no share and free to buy.
• The cooldown rule is one move: a sold state cannot buy, so the buy must start from rest.
• With no limit, three variables are enough, because each day only needs the values from the day before.
• With a limit k, keep one value per state for each count of buys, which is three arrays of length k + 1.
• The answer is the best value of rest or sold at the end, never hold.