[PYTHON] [Competition Pro] Summary of stock buying and selling problems

In real-world stock trading, we aim to earn more profits by buying new stocks or selling our own stocks while watching the ever-changing stock prices. Of course, it is very difficult to predict stock price movements, and this article does not aim for that either. In the actual situation, we do not know the future price movement, but this time we will simplify this and when given a column of numbers that record the price movement, when to "sell" or "buy" the highest profit Think about whether you can get **.

Problem setting

This is a problem of designing a function that takes an array containing the stock prices at each time as input and outputs the maximum profit obtained from it.

def calculate_max_profit(prices: List[int]) -> int:
  ...
  return max_profit

For example, in the graph above, the input is

prices = [7,1,5,3,6,4]

It looks like. The basic rules are as follows.

--You can either buy or sell at each time (you don't have to do anything, but you can't do both at the same time). --You can't sell before you buy --The buy operation and the sell operation must be alternated (buy → sell → buy → sell ...). ――If you buy a stock at a certain time, the stock price at that time will be deducted from your money, and if you sell the stock at a certain time, you will get the money for the stock price at that time. --The first money you have is 0. Your money may be 0 or less. ――I want to maximize the final amount of money I have.

If you know the price of each time in advance, it seems that you can calculate it immediately, but if the problem becomes complicated, it will not be so easy. There are various patterns depending on the conditions, so this article will introduce them all at once.

Case 1: If you can buy and sell only once

First, consider the most basic case where you can only buy and sell once. For example, if $ prices = [7,1,5,3,6,4] $, you can buy at time $ t = 1 $ and sell at time $ t = 4 $ as shown in the first figure. You can get a maximum profit of 5 (starting with $ t = 0 $). According to the rules, it is not possible to get a profit of 6 by first selling at $ t = 0 $ and buying at $ t = 1 $. On the other hand, if $ prices = [7,6,4,3,1] $, you can see that the profit is maximized at 0 if you never buy or sell.

The most basic solution is to search all the times to buy and sell.

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  max_profit = 0
  for i in range(n):  #Time to buy
    for j in range(i+1, n):  #Time to sell
      max_profit = max(max_profit, prices[j] - prices[i])  #Update maximum profit
                
  return max_profit

In this method, the time complexity of $ O (n ^ 2) $ is calculated with the length of $ prices $ as $ n $.

From here, we will consider a method to reduce the amount of time calculation. First, consider selling a stock at some point $ t = i $. If you decide to sell here, when is it desirable to buy stock? That is when the stock price is the cheapest before $ t = i $. In other words, if the minimum value of the stock price before $ t = i $ is known, the difference will be the maximum value of the profit that can be obtained. This means that you can calculate as long as you keep the minimum value while scanning $ prices $. 

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  min_value = float("inf")
  max_profit = 0
  for i in range(n):  #i is the time to sell
    max_profit = max(max_profit, prices[i] - min_value)
    min_value = min(min_value, prices[i])  #Update minimum
                
  return max_profit

This limits the time complexity to $ O (n) $ and the space complexity to $ O (1) $. Also, if the profit at each point is 0 or less, $ max \ _ profit = 0 $, and the result is that it is better not to buy or sell.

Case 2: If you can buy and sell up to 2 times

Next, let's change the number of times you can buy and sell up to 2 times. You can do "buy-> sell-> buy-> sell" in $ prices $. For example

prices = [3,3,5,0,0,3,1,4]

For, buy at $ t = 0 $ (-3) → sell at $ t = 2 $ (+5) → buy at $ t = 3 $ (-0) → sell at $ t = 7 $ (+) You can get a maximum profit of 6 by following the steps of 4) and. Increased complexity for one-time transactions. The simplest calculation method is to search the timing of each operation completely, but the amount of calculation is $ O (n ^ 4) $, which is not realistic.

Now, consider making the first "sell" at some point $ t = i $. You can find out when to make the first "buy" for this sale in the same way as in Case 1. So how can you calculate the maximum profit for the second buy and sell after this? You can do the same operation again for the part of $ prices $ after $ t = i + 1 $. The code would look like this:

def calculate_max_profit_twice(prices: List[int]) -> int:
  min_value = float("inf")
  max_profit = 0
  for i in range(n):
    current_profit = (prices[i] - min_value) + calculate_max_profit_once(prices[i+1:])
    max_profit = max(max_profit, current_profit)
    min_value = min(min_value, prices[i])
                
  return max_profit

Think of the $ calculate \ _max \ _profit \ _once () $ function as a Case 1 function. In this case, the amount of calculation could be suppressed to $ O (n ^ 2) $.

However, there are still useless operations. $ calculate \ _ max \ _profit \ _once () $ does the same calculation every time. For example, the calculation of the maximum profit after $ t = i $ and the calculation of the maximum profit after $ t = i + 1 $ should be performed in almost the same way.

Consider the case where you decide to "buy" a second stock at some point with $ t = j $ in order to reduce unnecessary calculations in this part. Then, when is the best time to sell this stock? Of course, it is the time when the stock becomes the highest after $ t = j $. This is the opposite of Case 1. In other words, by recording the maximum value after $ t = j $, the maximum profit when buying a stock at $ t = j $ can be obtained. So, specifically, first save the maximum profit of the second transaction after each $ j $ in the array $ second \ _ trans \ _max $. After that, find the maximum value of the first transaction at $ t = i $ as in Case 1, and add $ second \ _ trans \ _max [i + 1] $ to find the maximum profit for the second transaction. Can be done.

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  if n < 2:
    return 0
        
  #Calculate the maximum profit of the second transaction first
  second_trans_max = [0]*(n+1)
  max_after = float("-inf")
  for i in reversed(range(n)):
    second_trans_max[i] = max(max_after - prices[i], second_trans_max[i+1])  #Maximum profit of the second transaction
    max_after = max(max_after, prices[i])  #Update of maximum value after this
        
  #Calculate the maximum profit of the first transaction
  max_profit = first_trans_max = 0
  min_before = prices[0]
  for i in range(1, n):
    first_trans_max = max(prices[i]-min_before, first_trans_max)  #Maximum profit of the first transaction
    max_profit = max(first_trans_max + second_trans_max[i+1], max_profit)  #Overall maximum profit
    min_before = min(min_before, prices[i])  #Update of the previous minimum value
            
  return max_profit

In addition, this code includes the case where the number of transactions is 0 or 1. As a result, the time complexity can be reduced to $ O (n) $. On the other hand, the amount of spatial complexity is $ O (n) $ because it holds $ second \ _trans \ _max $.

Case 3: If you can buy and sell as many times as you like

Next, consider the case where you can buy and sell as many times as you like. I feel that the problem has become more complicated because I have to think about all the possibilities because the number of times is no longer limited, but in reality this setting is not difficult. All you have to do is repeat the operation of selling before the price is about to drop and buying at the cheapest price when the price is about to rise. 

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  valley = peak = prices[0]
  max_profit = 0
  idx = 0
        
  while idx < n-1:
    while idx < n-1 and prices[idx] >= prices[idx+1]:
      idx += 1
    valley = prices[idx]
    while idx < n-1 and prices[idx] <= prices[idx+1]:
      idx += 1
    peak = prices[idx]
    max_profit += peak - valley
                
  return max_profit

This code simply looks for and records the top and bottom of price movements. This is an intuitively valid strategy, but you can think of it as follows: For example, suppose you buy a stock at $ t = i $ and sell a stock at $ t = j $. If $ prices \ [a ]> prices \ [b ] ~~ (a <b) $ instead of a monotonous increase (in a broad sense) between $ i $ and $ j $, then $ You can make more profit by selling at t = a $ and buying at $ t = b $. In other words, you should always buy and sell at the point where the increase and decrease are switched. After all, this is the same as the calculation that "if the adjacent stock price is larger than the current time, it will be added to the profit, and if it is smaller, it will not be added". If you condense this, you can write the answer in one line.

def calculate_max_profit(prices: List[int]) -> int:
  return sum([max(a-b, 0) for a, b in zip(prices[1:], prices[:-1])]) if prices else 0

Here, the difference between adjacent values ​​is taken in order to add only the increase. In both codes, the time complexity is $ O (n) $ and the space complexity is $ O (1) $.

Case 4: If you can buy and sell up to k times

Next is the case where the upper limit of the number of buying and selling is set to $ k $. This is a general form of Case 1 and Case 2. However, if $ k $ becomes 3,4,5 ..., it cannot be calculated in the same way as Case 1 and 2. Also, unlike Case 3, it cannot be decided only by looking at the increase or decrease.

Here, the first thing I want to sort out before thinking about the solution is when $ k $ is large. At the extreme, if $ n = 10, k = 100 $, this is virtually the same situation as Case 3, where you can trade without limits. What is the line between Case 3 and Case 4? That is when $ k = \ frac {n} {2} $. The maximum number of transactions that can be traded in length $ n $ is $ \ frac {n} {2} $ times, so if $ k $ is larger than that, solving Case 3 is sufficient.

Now let's reconsider the case where $ k $ is not that big. For example, at a certain $ t = i $, the maximum profit $ max \ _sell [i] [j] $, $ j $ after finishing the $ j $ "sell" by $ t = i $. Suppose you know the maximum profit $ max \ _buy [i] [j] $ after you finish "buying" for each $ j $. Then, $ max \ _ sell $ and $ max \ _buy $ at $ t = i + 1 $ can be calculated as follows.

\begin{align}
max\_buy[i+1][j] &= \max(max\_buy[i][j], max\_sell[i][j-1] - prices[i+1]) \\
max\_sell[i+1][j] &= \max(max\_sell[i][j], max\_buy[i][j] + prices[i+1])
\end{align}

What does this mean? First, for $ max \ _buy $, the $ j $ th "buy" is done after the $ j-1 $ th "sell" according to the rules. Therefore, from the maximum profit of $ max \ _ sell [i] [j-1] $ at the end of the $ j-1 $ "sell" at the $ i $ point, at the $ i + 1 $ point. $ j $ Subtract $ prices [i + 1] $ for the third "buy". This is compared with the maximum profit $ max \ _buy [i] [j] $ at the end of the $ j $ th "buy" at the $ i $ point. Similarly, for $ max \ _sell $, the $ j $ "sell" is done after the $ j $ "buy", so the $ j $ "buy" at the $ i $ point is over. Add the maximum profit $ max \ _buy [i] [j] $ at the stage, and the profit $ prices [i + 1] $ obtained from the $ j $ th "sell" at the $ i + 1 $ point. I will. This is compared with the maximum profit $ max \ _sell [i] [j] $ at the end of the $ j $ th "sell" at the $ i $ point. This is exactly dynamic programming (DP) for $ i $ and $ j $. With proper initial and recurrence formulas, you can efficiently calculate values ​​for any $ i $ and $ j $.

def calculate_max_profit(prices: List[int], k: int) -> int:
  n = len(prices)
  if k == 0:
    return 0
        
  if k >= n//2:
    #For Case 3
    max_profit = 0
    for i in range(1, n):
      max_profit += max(0, prices[i] - prices[i-1])
    return max_profit
            
  else:
    #Initialize
    max_buy = [[float("-inf")]*(k+1) for _ in range(n)]
    max_sell = [[float("-inf")]*(k+1) for _ in range(n)]
    for i in range(n):
      max_sell[i][0] = max_sell[i][0] = 0
    max_buy[0][1] = -prices[0]

    # DP
    for i in range(1,n):
      for j in range(k):
        max_sell[i][j+1] = max(max_sell[i-1][j+1], max_buy[i-1][j+1]+prices[i])
        max_buy[i][j+1] = max(max_buy[i-1][j+1], max_sell[i-1][j]-prices[i])
                
    #Return the maximum profit of all j
    return max(max_sell[n-1])

The time complexity and spatial complexity of this are both $ O (nk) $. Also, by carefully considering the calculation order, the memory used for DP can be reduced.

def calculate_max_profit(prices: List[int]) -> int:
  ##Same as above before this
  else:
    #Initialize
    max_buy = [float("-inf")]*(k+1)
    max_sell = [float("-inf")]*(k+1)
    max_sell[0] = 0

    for i in range(n):
      for j in reversed(range(k)):
        max_sell[j+1] = max(max_sell[j+1], max_buy[j+1] + prices[i])
        max_buy[j+1] = max(max_buy[j+1], max_sell[j] - prices[i])

    return max(max_sell)

What they are doing is the same, but by overwriting $ max \ _ sell $ and $ max \ _buy $, the amount of spatial complexity can be reduced to $ O (k) $.

* Solving other problems with DP

By the way, this solution solves the problem by recording "the maximum profit in each state after" selling "and" buying "at the $ i $ point ($ k $ th time)". Let's use this idea to solve the problems so far. Case 2 is a special case where $ k = 2 $ in Case 4, so it can be solved in the same way.

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  max_buy1 = max_buy2 = float("-inf")
  max_sell1 = max_sell2 = 0
  for i in range(n):                             
    max_sell2 = max(max_sell2, max_buy2 + prices[i])
    max_buy2 = max(max_buy2, max_sell1 - prices[i])
    max_sell1 = max(max_sell1, max_buy1 + prices[i])
    max_buy1 = max(max_buy1, -prices[i])
        
  return max(sell1, sell2)

Only 4 variables are needed, and the amount of spatial complexity is $ O (1) $. Case 3 can also be solved as follows.

def calculate_max_profit(prices: List[int]) -> int:
  max_sell, max_buy = 0, float("-inf")
  for i in range(n):
    sell, buy = max(sell, buy+prices[i]), max(buy, sell-prices[i])

  return max_sell

Since you can trade as many times as you like, you don't have to think about $ k $, and $ sell $ and $ buy $ are calculated at the same time.

Case 5: When there is a charge for buying and selling

Case 4 was the main theme of this article, but I would like to introduce two other settings for the slightly changing sphere. The first is the case where "a fee is charged for one transaction (buying and selling)". In this case, you can freely trade as many times as you like, but if the profit you get is small, you will incur a loss due to the commission. How should I think about it?

In fact, this is easy, just subtract $ fee $ when updating $ max \ _sell $ in Case 3. $ max \ _ buy + prices [i] --You can decide whether to trade according to the size of fee $.

def calculate_max_profit(prices: List[int], fee: int) -> int:
  n = len(prices)
  max_sell, max_buy = 0, float("-inf")
  for i in range(n):
    max_sell, max_buy = max(max_sell, max_buy + prices[i] - fee), max(max_buy, max_sell - prices[i])
                
  return max_sell

Case 6: When you cannot buy or sell in a row

Finally, consider the condition that "at least one dose must be inserted between the sale of the stock and the next purchase." In this case, add the state $ max \ _ stay $ that represents a break to Case 3. And if the update formula is as follows, it is possible to express the state of having a break.

\begin{align}
max\_sell[i+1] &= \max(max\_sell[i], max\_buy[i] + prices[i+1]) \\
max\_buy[i+1] &= \max(max\_buy[i], max\_stay[i] - prices[i+1]) \\
max\_stay[i+1] &= \max(max\_stay[i], max\_sell[i])
\end{align}

max \ _sell [i + 1] is the maximum profit of "buy" up to $ i $ point $ max \ _buy [i] $ plus the profit of "sell" at $ i + 1 $ point $ prices [i + 1] ] $ Added, max \ _buy [i + 1] is the maximum profit after a break to the $ i $ point $ max \ _ stay [i] $ to $ i + 1 $ point "buy" The cost of $ prices [i + 1] $ minus, max \ _stay [i + 1] is the maximum profit of "selling" up to the $ i $ point $ max \ _ sell [i] $ in a resting state It is a transition. If you write this in a form that saves memory, it will be as follows.

def calculate_max_profit(prices: List[int]) -> int:
  n = len(prices)
  max_stay, max_sell, max_buy = 0, 0, -float("inf")
  for i in range(n):
    max_stay, max_sell, max_buy = max(max_stay, max_sell), max(max_sell, max_buy + prices[i]), max(max_buy, max_stay - prices[i])
            
  return max(max_stay, max_sell)

Summary

So far, we have comprehensively looked at stock price buying and selling problems, but by calculating the state representing ** "maximum profit of $ k $ buying and selling up to $ i $ at each coordinate" **, many problems can be solved. I found that I could handle it. If you read this article, you should be able to get the maximum profit by trading stocks ** if you have a clear future. You can actually run the code on the following sites that I referred to, so if you are interested, please try it!

** Reference (from LetCode) ** Best Time to Buy and Sell Stock Best Time to Buy and Sell Stock II Best Time to Buy and Sell Stock III Best Time to Buy and Sell Stock IV Best Time to Buy and Sell Stock with Cooldown Best Time to Buy and Sell Stock with Transaction Fee

(Addition) Supplement: Solving O (n log n)

In the comments, he pointed out that there is a more optimal method. Since you gave a detailed explanation in the comment section, please refer to that for details, but here, based on that, we will implement an implementation with a time complexity of $ O (n \ log n) $ for Case 4. I will put it.

def calculate_max_profit(prices: List[int], k: int) -> int:
    prices = [float("inf")] + prices + [float("-inf")]  #Insert infinite values ​​at both ends
    dif_array = [prices[i] - prices[i-1] for i in range(1, len(prices))]  #An array that stores the difference from the neighbor

    comp_dif = [dif_array[0]]  #Compressed array
    for i in range(1, len(dif_array)):
        if comp_dif[-1]*dif_array[i] < 0:
            comp_dif.append(dif_array[i])
        else:
            comp_dif[-1] += dif_array[i]

    max_profit = sum(comp_dif[1::2])  #Maximum profit when you can trade as many times as you want

    max_trans = (len(comp_dif)-1)//2  #Number of transactions to get maximum profit

    if max_trans < k:
        return max_profit

    value = copy.deepcopy(comp_dif)  #An array that stores the combined value
    forwardLength = [1]*len(comp_dif)  #An array showing how many are connected backward
    backwardLength = [1]*len(comp_dif)  #An array showing how many are connected forward

    que = [(abs(comp_dif[i]), i, forwardLength[i]) for i in range(len(comp_dif))]  #Initialize the heap with the absolute value of each value in the array after compression
    heapq.heapify(que)

    num = max_trans - k
    while que and num:  #Max for too many transactions_Deduct value from profit
        val, i, fi = heapq.heappop(que)
        if fi != forwardLength[i]:
            continue

        start = i - backwardLength[i-1]  #The beginning of the array to join
        mid = i + forwardLength[i]  #The beginning of the second half of the array to join
        end = i + forwardLength[i] + forwardLength[i + forwardLength[i]] - 1  #The end of the array to join

        forwardLength[start] = backwardLength[end] = end - start + 1
        forwardLength[i] = forwardLength[mid] = 0

        value[start] = value[start]+value[i]+value[mid]
        heapq.heappush(que, (abs(value[start]), start, forwardLength[start]))

        max_profit -= val
        num -= 1

    return max_profit

Recommended Posts

[Competition Pro] Summary of stock buying and selling problems
Problems of liars and honesty
Problems of liars and honesty
[For competition professionals] Summary of doubling
Summary of Python indexes and slices
Summary of modules and classes in Python-TensorFlow2-
Summary of OSS tools and libraries created in 2016
Summary of the differences between PHP and Python
Installation of Python3 and Flask [Environment construction summary]
Confirmation of basics of competition professionals ~ gcd and lcm ~
I / O related summary of python and fortran
About problems and solutions of OpenPyXL (Ver 3.0 version)
Summary of pickle and unpickle processing of user-defined class