Best Trade

best_trade.py

"""
Best Trade

Given a list of prices of some stock, determine the maximum
money that could be made from buying and then selling once.
That is, find the value (prices[sell_index] - prices[buy_index])
such that len(prices) > sell_index >= buy_index >= 0 and for all
other such index pairs (sell_index_2, buy_index_2), it is
the case that:

    prices[sell_index] - prices[buy_index] >= prices[sell_index_2] - prices[buy_index_2]
"""


def best_trade(prices):
    # O(1) time, O(1) space
    lowest_price_so_far = float('inf')
    best_trade_so_far = 0

    for p in prices:
        lowest_price_so_far = min(p, lowest_price_so_far)
        best_trade_so_far = max(best_trade_so_far, p - lowest_price_so_far)

    return best_trade_so_far


def best_trade_indices(prices):
    "Return the earliest pair of indices that get you the best trade."
    # O(1) time, O(1) space
    best_trade_so_far = 0
    lowest_price_index_so_far = buy_index = sell_index = 0

    for i,p in enumerate(prices):
        if p < prices[lowest_price_index_so_far]:
            lowest_price_index_so_far = i
        if p - prices[lowest_price_index_so_far] > best_trade_so_far:
            best_trade_so_far = p - prices[lowest_price_index_so_far]
            buy_index, sell_index = lowest_price_index_so_far, i

    return buy_index, sell_index


def tests():
    def brute_force(prices):
        # O(N^2) time, O(1) space
        if not prices:
            return 0

        trades = (prices[j] - prices[i]
                  for i in xrange(len(prices))
                  for j in xrange(i, len(prices)))

        return max(trades)

    def brute_force_indices(prices):
        # O(N^2) time, O(1) space
        if not prices:
            return (0, 0)

        index_pairs = ((i, j)
                       for i in xrange(len(prices))
                       for j in xrange(i, len(prices)))

        return max(index_pairs, key=lambda pair: prices[pair[1]] - prices[pair[0]])

    hard_coded_cases = (
        ([], 0, 0, 0),
        ([1], 0, 0, 0),
        ([3,2,1], 0, 0, 0),
        ([1,2,3], 2, 0, 2),
        ([10,11,12,1,2,4,3], 3, 3, 5),
        ([10,11,12,1,2,15,3], 14, 3, 5)
    )

    for prices,expected_best_trade,buy_index,sell_index in hard_coded_cases:
        assert best_trade(prices) == brute_force(prices) == expected_best_trade
        assert best_trade_indices(prices) == brute_force_indices(prices) == (buy_index, sell_index)

    import random

    for _ in xrange(1000):
        n = random.randint(2, 400)
        prices = [random.randint(0, 10000) for _ in xrange(n)]
        assert best_trade(prices) == brute_force(prices)
        assert best_trade_indices(prices) == brute_force_indices(prices)

    print 'Tests pass!'