【LeetCode OJ】Best Time to Buy and Sell Stock III

Problem Link:

http://oj.leetcode.com/problems/best-time-to-buy-and-sell-stock-iii/

---

Linear Time Solution

We try to solve this problem in O(n) time in the help of the algorithm in Best Time to Buy and Sell Stock, which can return the max profit by given a list of prices.

As transaction definition in Best Time to Buy and Sell Stock, we modify the algorithm that returns the best transaction (b, s) by given price list prices[0:n] where

1. prices[i] <= prices[s0], for i = b0, ..., n-1
2. prices[i] >= prices[b0], for i = 0, ..., s0

By the help of method find_best_transaction, the following algorithm will solve the problem in O(n) time:

1. By given the price list prices[0..n01], find the best transaction (b,s) = find_best_transaction(prices)
2. Find the maximum profit for the following three parts divided by b and s:
    P1: the maximum profit for the price list price[0..b-1]
    P2: the maximum profit for the price list price[b+1..s-1].inverse()
    P3: the maximum profit for the price list price[s+1..n-1]
3. Return prices[s]-prices[b] + max(P1,P2,P3)
 

---

Correctness

Now we woul prove the algorithm above is correct. Sicne we are asked to find at most two transactions, there are two cases: 1) (b0, s0) is one of the two transactions; 2) (b0, s0) is not in the result. For case 1) we can solve the problem by find the max profit for prices[0:b0] and prices[s0:n], and choose the bigger one as the second transaction.

Now we consider the case 2), let (b1, s1) and (b2, s2) be the two transactions where 0 <= b1 < s1 < b2 < s2 <= n-1. Also, we can have the following corolarries:

1. b0 <= s1 <= s0, we prove it by contradiction:
   1. if s1 < b0 (-----s1--b0-----s0------), then (b2,s2) could be replaced by (b0, s0).
   2. if s1 > s0 (----b0----s0--s1--------), then (b1, s1) could be replaced by (b0, s0).
2. b0 <= b2 <= s0, we prove it by contradiction:
   1. if b2 < b0 (---b2----b0-------s0---), then (b2, s2) could be replaced by (b0, s0).
   2. if b2 > s0 (---------b0----s0--b2--), then (b1, s1) could be replaced by (b0, s0).
3. For b1 < s1, since s1 <= s0, then s1 could be b0, since prices[b0] <= prices[i] for i = 0,...,s0.
4. For s2 > b2, since b2 >= b0, then s2 could be s0, since prices[s0] <= prices[i] for i = b0, ..., n-1

Therefore, we can find the best s1 and b2 scanning between b1(b0) and s2(s0), which only takes O(n) time and is a little tricky.

Suppose the optimal trasactions can be as follows:

　　　　-----b1(b0)----s1------b2------s2(s0)----

The profit should be (p[s1] - p[b1]) + (p[s2]-p[b2]) which can be rewritten as (p[s2]-p[b1]) + (p[s1]-p[b2]). Therefore, it equals to finding best transaction (b2, s1) where the given prices list is the inverse of prices[b1+1, ..., s2-1].

From the two cases above, our algorithm will give the correct answer.

---

Python Code

The following is the python implementation accepted by oj.leetcode.com.

class Solution:
    # @param prices, a list of integer
    # @return an integer
    def maxProfit(self, prices):
        n = len(prices)
        if n < 2:
            return 0
        # Find the single best transaction
        b0, s0, profit0 = self.find_best_transaction(prices)

        # Calculate the max profit of the three parts partitioned by (b0, s0)
        profit1 = profit2 = profit3 = 0
        if b0 > 0:
            profit1 = self.find_best_transaction(prices[:b0])[2]
        if s0 > b0 + 1:
            profit2 = self.find_best_transaction(prices[b0+1:s0][::-1])[2]
        if s0 < n-1:
            profit3 = self.find_best_transaction(prices[s0+1:])[2]

        # Return the best case
        return profit0 + max(profit1, profit2, profit3)

    def find_best_transaction(self, prices):
        """
        Return the transaction (b,s) that obtains maximum profit.
        @param prices: a list of prices
        @return: (b,s) where 0 <= b < s < len(prices)
        """
        # Initialize with prices[0]
        b = s = l = 0
        # Scan from i = 1 to n-1
        for i in xrange(1, len(prices)):
            if prices[i] <= prices[l]:
                l = i
            elif prices[i] > prices[s] or prices[s] - prices[b] < prices[i] - prices[l]:
                s = i
                b = l
        return b, s, prices[s]-prices[b]