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