Pythonで解く「株式売買のベストタイミング III」― 最大2回の取引で最大利益を求める方法

ある配列が与えられ、その i 番目の要素が i 日目の株価を表しているとします。このとき、最大2回までの取引で得られる最大の利益を求めるアルゴリズムを設計するのが、この問題の目的です。

例として、株価が [3,3,5,0,1,3,1,4] の場合を考えてみましょう。4日目に価格 0 で株を買い、6日目に価格 3 で売れば利益は 3 − 0 = 3。その後、7日目に価格 1 で買い、8日目に価格 4 で売れば利益は 4 − 1 = 3 となります。合計の利益は 3 + 3 = 6 であり、これがこのケースでの最大値です。

解法のアプローチ

この問題は動的計画法(DP)を用いることで効率よく解けます。鍵となるのは、計算を「前方向」と「後ろ方向」の2つのパスに分けることです。

• 前方向のパス: 各日 i において、それまでの最小購入価格 xmin を追跡しながら、dp[i](i 日目までに1回の取引で得られる最大利益)を求めます。
• 後ろ方向のパス: 各日 i 以降の最高価格 xmax[i] を求め、それをもとに「i 日目より後に行う1回の取引の最大利益」tempp を計算します。
• 統合: dp[i](前半の取引)+ tempp(後半の取引)の組み合わせの中で最大になるものが、2回の取引で得られる最大利益となります。

手順の詳細

• p が空なら 0 を返します。
• n := p の要素数とし、長さ n の dp 配列、ans := 0、xmin := p[0] で初期化します。
• i を 1 から n − 1 まで順に処理します。
  • xmin := min(xmin, p[i])
  • dp[i] := max(dp[i], p[i] − xmin)
  • ans := max(ans, dp[i])
• 長さ n の xmax 配列を作り、xmax[n−1] := p[n−1] とします。
• i を n − 2 から 0 まで逆順にたどりながら、xmax[i] := max(xmax[i+1], p[i]) を計算します。
• tempp := 0 とし、再度 i を n − 2 から 0 まで逆順にたどります。
  • tempp := max(tempp, xmax[i+1] − p[i+1])(i 日目以降に行う1回の取引の最大利益)
  • ans := max(ans, dp[i] + tempp)(前半と後半の取引を組み合わせる)
• 最後に ans を返します。

この手法により、時間計算量 O(n)・空間計算量 O(n) で問題を解くことができます。

実装例

以下はPythonによる実装です。コードを追いかけることで、仕組みへの理解がさらに深まります。

class Solution(object):
   def maxProfit(self, p):
      if not p:
         return 0
      n = len(p)
      dp = [0 for i in range(n)]
      ans = 0
      xmin = p[0]
      for i in range(1,n):
         xmin = min(xmin,p[i])
         dp[i] = max(dp[i],p[i]-xmin)
         ans = max(ans,dp[i])
      xmax = [0 for i in range(n)]
      xmax[-1] = p[-1]
      tempp = 0
      for i in range(n-2,-1,-1):
         xmax[i] = max(xmax[i+1],p[i])
      for i in range(n-2,-1,-1):
         tempp = max(tempp,xmax[i+1]-p[i+1])
         ans = max(ans,dp[i]+tempp)
      return ans
ob = Solution()
print(ob.maxProfit([3,3,5,0,1,3,1,4]))

入力

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

出力

6