Add investment_optimization.py exercise
This commit is contained in:
parent
61d3c542d1
commit
1b4ea820a0
168
dotfiles/lib/python/investment_optimization.py
Executable file
168
dotfiles/lib/python/investment_optimization.py
Executable file
@ -0,0 +1,168 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
import heapq
|
||||||
|
from collections import namedtuple
|
||||||
|
|
||||||
|
|
||||||
|
BuyOpportunity = namedtuple(
|
||||||
|
'BuyOpportunity',
|
||||||
|
['trans_yield', 'start', 'end']
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def maximize_profit(num_transactions, prices):
|
||||||
|
minima, maxima = find_extrema(prices)
|
||||||
|
opp_queue, reverse_opp_queue = make_opportunity_queues(
|
||||||
|
minima, maxima, prices,
|
||||||
|
)
|
||||||
|
if not opp_queue:
|
||||||
|
return []
|
||||||
|
|
||||||
|
largest_segment = opp_queue[0][1]
|
||||||
|
# Segments will be kept in sorted order
|
||||||
|
segments = [largest_segment]
|
||||||
|
|
||||||
|
# Remove any reverse yields that are greater than the largest actual yield
|
||||||
|
# since they can never be realized anyway.
|
||||||
|
while (reverse_opp_queue and reverse_opp_queue[0][1].trans_yield >=
|
||||||
|
largest_segment.trans_yield):
|
||||||
|
heapq.heappop(reverse_opp_queue)
|
||||||
|
|
||||||
|
def try_rev_opp():
|
||||||
|
# It is okay to definitely pop here even though we don't know that we
|
||||||
|
# can actually use the opp for the following reason:
|
||||||
|
# Since the rev opp queue was selected OVER that of the opp queue, we
|
||||||
|
# KNOW that the bounding segment that includes this rev opp must have
|
||||||
|
# already been selected if it is going to be included at all (since it
|
||||||
|
# must have greater yield).
|
||||||
|
_, rev_opp_can = heapq.heappop(reverse_opp_queue)
|
||||||
|
for (seg_index, split_seg) in enumerate(segments):
|
||||||
|
if split_seg.end >= rev_opp_can.end:
|
||||||
|
# Since segments is sorted, this must be the correct segment
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
return
|
||||||
|
if split_seg.start <= rev_opp_can.start:
|
||||||
|
# We found the containing segment
|
||||||
|
left_yield = prices[rev_opp_can.start] - prices[split_seg.start]
|
||||||
|
right_yield = prices[split_seg.end] - prices[rev_opp_can.end]
|
||||||
|
left_segment = BuyOpportunity(left_yield, split_seg.start, rev_opp_can.start)
|
||||||
|
right_segment = BuyOpportunity(right_yield, rev_opp_can.end, split_seg.end)
|
||||||
|
segments.pop(seg_index)
|
||||||
|
segments.insert(seg_index, left_segment)
|
||||||
|
segments.insert(seg_index + 1, right_segment)
|
||||||
|
|
||||||
|
def try_opp():
|
||||||
|
_, opp = heapq.heappop(opp_queue)
|
||||||
|
if not segments:
|
||||||
|
segments.append(opp)
|
||||||
|
insertion_index = 0
|
||||||
|
for (index, seg) in enumerate(segments):
|
||||||
|
if seg.start >= opp.start:
|
||||||
|
insertion_index = index
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
insertion_index = len(segments)
|
||||||
|
seg = None
|
||||||
|
previous_seg = segments[insertion_index - 1] if insertion_index > 0 else None
|
||||||
|
|
||||||
|
if ((seg is None or seg.start >= opp.end) and
|
||||||
|
(previous_seg is None or previous_seg.end <= opp.start)):
|
||||||
|
# There is no overlap, so we can insert
|
||||||
|
segments.insert(insertion_index, opp)
|
||||||
|
else:
|
||||||
|
pass
|
||||||
|
|
||||||
|
while (opp_queue or reverse_opp_queue) and len(segments) < num_transactions:
|
||||||
|
if not reverse_opp_queue:
|
||||||
|
try_opp()
|
||||||
|
elif not opp_queue:
|
||||||
|
try_rev_opp()
|
||||||
|
else:
|
||||||
|
opp_can = opp_queue[0][1]
|
||||||
|
rev_opp_can = reverse_opp_queue[0][1]
|
||||||
|
|
||||||
|
if rev_opp_can.trans_yield > opp_can.trans_yield:
|
||||||
|
try_rev_opp()
|
||||||
|
else:
|
||||||
|
try_opp()
|
||||||
|
|
||||||
|
return segments
|
||||||
|
|
||||||
|
|
||||||
|
def make_opportunity_queues(minima, maxima, prices):
|
||||||
|
opp_queue = []
|
||||||
|
reverse_opp_queue = []
|
||||||
|
for min_index, minimum in enumerate(minima):
|
||||||
|
for max_index, maximum in enumerate(maxima):
|
||||||
|
transaction_yield = prices[maximum] - prices[minimum]
|
||||||
|
if transaction_yield < 0:
|
||||||
|
# We can ignore this pair because the transaction has negative
|
||||||
|
# yield.
|
||||||
|
continue
|
||||||
|
# minimum comes before maximum in time
|
||||||
|
if minimum < maximum:
|
||||||
|
# Transaction yield is made negative because heapq is a min-heap
|
||||||
|
heapq.heappush(
|
||||||
|
opp_queue, ((-transaction_yield, maximum - minimum), BuyOpportunity(
|
||||||
|
transaction_yield, minimum, maximum,
|
||||||
|
)),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
heapq.heappush(
|
||||||
|
reverse_opp_queue, (-transaction_yield, BuyOpportunity(
|
||||||
|
transaction_yield, maximum, minimum,
|
||||||
|
))
|
||||||
|
)
|
||||||
|
return opp_queue, reverse_opp_queue
|
||||||
|
|
||||||
|
|
||||||
|
def find_extrema(prices):
|
||||||
|
maxima = []
|
||||||
|
minima = []
|
||||||
|
length_of_prices = len(prices)
|
||||||
|
if length_of_prices < 2:
|
||||||
|
return minima, maxima
|
||||||
|
|
||||||
|
upwards = None
|
||||||
|
last = prices[0]
|
||||||
|
|
||||||
|
for (index, price) in enumerate(prices):
|
||||||
|
if price < last:
|
||||||
|
if upwards is True:
|
||||||
|
maxima.append(index - 1)
|
||||||
|
elif upwards is None:
|
||||||
|
# We set the starting price as a maximum, but theres no point
|
||||||
|
# since we would really never buy.
|
||||||
|
maxima.append(0)
|
||||||
|
pass
|
||||||
|
upwards = False
|
||||||
|
elif price > last:
|
||||||
|
if upwards is False:
|
||||||
|
minima.append(index - 1)
|
||||||
|
elif upwards is None:
|
||||||
|
# The starting value is a minimum
|
||||||
|
minima.append(0)
|
||||||
|
upwards = True
|
||||||
|
last = price
|
||||||
|
|
||||||
|
if upwards is True:
|
||||||
|
maxima.append(length_of_prices - 1)
|
||||||
|
elif upwards is False:
|
||||||
|
minima.append(length_of_prices - 1)
|
||||||
|
|
||||||
|
return minima, maxima
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
print (maximize_profit(10, [0, 1, 3, 2, 3, 0, 10, 12, 1, 2, 3, 2, 0, 2, 4, 3, 6, 4, 14, 1, 0, 2, 4, 5, 4, 5, 6]))
|
||||||
|
|
||||||
|
print [
|
||||||
|
BuyOpportunity(trans_yield=1, start=0, end=1),
|
||||||
|
BuyOpportunity(trans_yield=12, start=2, end=4),
|
||||||
|
BuyOpportunity(trans_yield=2, start=5, end=7),
|
||||||
|
BuyOpportunity(trans_yield=6, start=9, end=11),
|
||||||
|
BuyOpportunity(trans_yield=6, start=12, end=13),
|
||||||
|
BuyOpportunity(trans_yield=10, start=14, end=15),
|
||||||
|
BuyOpportunity(trans_yield=5, start=17, end=20),
|
||||||
|
BuyOpportunity(trans_yield=2, start=21, end=23),
|
||||||
|
]
|
Loading…
Reference in New Issue
Block a user