Best time to buy and sell stocks when allowing consecutive buys or sells

Question

Problem

You're given the n stock prices for n days. Output the maximum profit you can reach by trading stocks. You can only trade at most once a day: on each day you can choose to either buy a single stock, or sell a single stock (if you have one), or give up the trade for that day and do nothing.

Example 1:

Given a = [1,2,10,9], return 16

Explanation:

You can buy on day 1 and 2 and sell on day 3 and 4.

Profit: -1-2+10+9 = 16

Example 2:

Given a = [9,5,9,10,5], return 5

Explanation:

You can buy on day 2 and sell on day 4.

Profit: -5 + 10 = 5

My analysis

The difficult part is that you can engage in consecutive buys and/or sells, meaning that once you posses a stock, you don't necessarily have to sell it before buying another one.

My idea is the following algorithm:

Start from the largest price, and then match the smallest price that occurs before that maximum price in the input array. After matching, remove these two prices from the array and keep repeating this process until you can find no more match. It seems like this algorithm works, but it costs O(n²) time, which is not fast enough.

Question

How could this be solved with a better time complexity, such as O(nlogn)?

Answer 1

We can model this as a min-cost circulation problem and solve it optimally with a specialized O(n log n)-time algorithm similar to your idea.

In the flow network, there is a node for each day and a node representing the market. There are two unit-capacity arcs for each day, one from the market with cost equal to the price that day, one to the market with cost equal to minus the price. There are arcs of zero cost and unbounded capacity that can move flow from each day (except the last) to the one after it. These represent holding stock.

Using () to represent nodes, ==> to represent unbounded capacity arcs and --> to represent unit capacity arcs, and labeling the costs, your sample instance is

      0        0        0
 ()======>()======>()======>()
 ^\       ^\       ^\       ^\
1| |-1   2| |-2  10| |-10  9| |-9
  \v       \v       \v       \v
  (                            )

Technically it's possible in this reformulation to both buy and sell on the same day, but that's not a profitable move, so it doesn't matter.

Given a residual network, the theory (linear programming duality) says we're done if and only if there's no negative-cost simple cycle. The intuitive meaning of such cycles is exactly what you would expect: buying a share and selling it profitably later.

The algorithm works by successively eliminating all negative-cost simple cycles (profitable cycles from now on) on the first k days for k from 1 to n. In the base case k = 1, the first day alone is never profitable, so we can move along to the inductive step.

For the inductive step, we know that there are no profitable cycles in the first k-1 days and want to extend that to k. If there is a profitable cycle in the first k days, it involves selling on day k. But what to buy? We can answer that question efficiently by maintaining a min-priority queue of our residual buying opportunities. We compare the day k price to the queue min, and if it's higher we do the deal, which involves popping the min and pushing the day k price, since from the point of view of the residual network, canceling our sale later looks the same as buying a share. Then we push the day k price regardless to represent the possibility of actually buying on day k.

We have to be careful here and prove that we didn't just introduce another profitable cycle. That's the reason to choose the min: we can't combine the new "selling" (actually canceling the buy) opportunity profitably with any of the residual buying opportunities, because the new selling price was not greater than any of those opportunities.

The finished algorithm is pretty simple. In Python:

import heapq


def trading_profit(prices):
    profit = 0
    queue = []
    for price in prices:
        if queue and queue[0] < price:
            profit += price - queue[0]
            heapq.heapreplace(queue, price)
        heapq.heappush(queue, price)
    return profit

Answer 2

This is an O(n²) algorithm. So in that sense it does not answer your question for something asymptotically faster, but as in a comment you learned that you're algorithm won't work, I believe it may be useful nonetheless.

I'd go for dynamic programming. Iterate over the days, and maintain a list where the index describes the number of stock you have, and the value is the best cash balance to arrive in that situation. So start with the list being [0], i.e. a single entry indicating that you can have zero stock at balance zero.

For each day, you can buy, sell or skip. You can express all the together using something like this:

balance_new[i] = max(balance[i], balance[i-1] - quote, balance[i+1] + quote)

The first entry represents skip: you keep current stock and balance. The second entry represents buy: you gain one stock (from i-1 to i) but reduce balance by the day's price. The third entry is a sell: you reduce stock by one but gain the current price to your balance.

The balance_new you get from this becomes the balance for the next day. And you'll need to take some care around the boundary of the list, where one of the expressions becomes invalid because it would index out of bounds. You can't get to zero stock with a buy operation. The requested maximum profit is balance[0] after you've processed all days. It represents the maximum balance that leaves you with no stock.

You have an outer loop iterating over n days. And an inner loop iterating over the potential number of stock you might possess at that point. That number grows linear in each iteration. If you want you can try to be clever and reduce the number of steps for the inner loop by one after you reached half the steps of the outer loop. That's because it never pays to acquire more stock than you can sell by the end. So the number of steps in the inner loop would go from something like one to something like n/2 then back down again, for a total of n²/4 + O(n) but that's still in O(n²) over all.