Mathwizurd.com is created by David Witten, a mathematics and computer science student at Vanderbilt University. For more information, see the "About" page.

Buying/Selling Stocks

This will be a series of problems, getting more and more difficult.

Problem 1: Infinite buying/selling

Let’s say you have a list of stock prices, which reflect the price of the stock on a given day:

[7,1,5,3,6,8,4]

What’s the maximum profit you can get from buying/selling stocks, completing as many transactions as you like (i.e., buy one and sell one share of the stock multiple times)?

Answer: 9 (1 -> 5, 3 -> 8)

Note: You may not engage in multiple transactions at the same time (i.e., you must sell the stock before you buy again).

Solution:

For me, the way I thought about this was that every long-term transaction can be broken into smaller transactions. Buying on day 4 and selling on day 6 is equivalent to buying on day 4 and selling on day 5, then buying on day 5 and selling on day 6. What determines whether you should buy on a given day? If the next day is greater!

So, now we get our very simple solution: We loop through every day, and we buy whenever the next day is greater than the current day.

def maxProfit(array):
    total = 0
    for i in range(len(array) - 1):
        if array[i] < array[i + 1]:
            total += array[i-1] - array[i]
    return total

Problem 2: Buy/sell once

This is exactly the same problem, except you are only allowed to make one transaction the whole time.

Solution:

If you’re asked this in an interview, the safest thing to do is to think of a brute force solution.

O(n^2): Brute Force

The most basic solution is to try out every possible combination of buying and selling, and keep track of the best result.

E.g.: [7,1,5,3,6,8,4] → 5

[7,1,5,3,6,8,4] →3

There are n*(n-1)/2 total operations, or O(n) time.

O(n log n): Sorting

There is no way to sort here, because the order of the prices matters.

O(n): Single Pass

Based on the logic above, we’re not gonna see an interview problem with a brute force solution as the ideal solution (usually). So, you can deduce that this is the complexity of the ideal solution.

Because it’s a single pass, we can try selling at every location, and finding the maximum of that. The question is when would you buy it if you sell a stock today? You would buy it at the minimum up until that point. So, that’s all you have to do:

At every point, keep track of the minimum, and calculate current price - minimum. Take the maximum out of all of those.

def maxProfit(array):
    if len(array) < 2:
        return 0
    minimum, best = array[0], 0
    for i in array[1:]:
        if i < minimum:
            minimum = i
        elif i - minimum > best:
            best = i - minimum  
    return total

Problem 3: Infinite buy/sell with cooldown

In this problem, you can buy/sell stocks as many times as you want. However, after you sell your stock, you cannot buy stock on next day.

On any given day, you have two options: you can either buy or sell stock. Both of them come with preconditions. First, in order to sell, the last transaction you’ve made was a sale. In order to buy, the last transaction you’ve made was a sale 2 days ago (or it’s your first time buying).

On every day, you have two numbers to keep track of. The highest amount you could’ve made if your last transaction was a sale, call it sell[i], and the highest amount you could’ve made if your last transaction was a buy, call it buy[i].

buy[i] = max(buy[i-1], sell[i-2] - prices[i])

Your last action was either already a buy, in which case you can’t buy again (buy[i-1]), or your last action was a sale two days ago, and you’re buying again today.
sell[i] = max(sell[i-1], buy[i-1] + prices[i])

Your last action was either already a sale (sell[i-1]), or your last action was buying, and you’re selling today.

That’s it! All we have to do now is pick a starting condition. buy[i] means your last transaction was a buy, so on the first day, buy[0] means we bought the stock, which means buy[0] = -1 * prices[0]. Likewise, sell[0] = 0, because we’re selling nothing.

Solution

def maxProfit(prices):
    buy = [0 for i in prices]
    sell = [0 for i in prices]
    buy[0] = -1 * prices[0]
    for i in range(1, len(prices)):
        if i < 2:
            buy[i] = max(buy[i-1], prices[i])
        else:
            buy[i] = max(buy[i-1], sell[i-2] + prices[i])
        sell[i] = max(sell[i-1], buy[i-1] + prices[i])
    return sell[-1] #We're not gonna end on a purchase!

Cleaner Solution

Timewise, this is an amazing solution. It’s O(n) with one pass through the array. However, space-wise, it's inefficient. You don’t need to save your buy/sell information from 5 days ago! All you need is buy[i-1], sell[i-1], and sell[i-2]. Let’s rename these:

We can just refer to buy[i-1] as buy. After all, it’s just whatever information was previously stored in buy.
We can refer to sell[i-1] as sell.
We can refer to sell[i-2] as rest.

Now, let’s rewrite this code:
It’s important to note that doing something like a,b,c = 1,2,3 means that they happen simultaneously, so you don’t have to worry about the order you write the expressions in.

def maxProfit(prices):
    buy, sell, rest = -1 * prices[0],0,0
    for price in prices[1:]:
        buy,sell,rest = max(buy, rest - price),\
          max(sell, buy + price), sell
    return sell 

Problem 4: Infinite buy/sell with a transaction fee

This is extremely similar to the previous problem. You have to keep track of two things: the maximum amount you can make if your last action was a buy, and the maximum amount you can make if your last action was a sale.

buy[i] = max(buy[i-1], sell[i-1] - prices[i])

This is exactly the same as before, except we don’t care about the cool-down period.

sell[i] = max(sell[i-1], buy[i-1] + prices[i] - fee)

Either you sold your stock earlier or you sell it today and pay a fee.

Right off the bat, we know we can simplify it. We only need to keep track of two variables: sell[i-1] and buy[i-1]. We know our initial conditions: buy = -1 * prices[0], and sell = 0

def maxProfit(prices, fee):
    buy, sell = -1 * prices[0], 0
    for price in prices[1:]:
        buy, sell = max(buy, sell - price),\
          max(sell, buy + price - fee)
    return sell  # We don't want to end on a purchase

Maximum Subarray