Pythonで配列を分割して合計を最大化する方法(動的計画法)

問題概要

整数配列 A が与えられたとき、この配列を「長さが K 以下の連続する部分配列」に分割することを考えます。分割後、各部分配列に含まれるすべての要素は、その部分配列内の最大値に置き換えられます。求めたいのは、分割後の配列の合計値として考えられる最大値です。

例えば、入力が [1, 15, 7, 9, 2, 5, 10]、K = 3 の場合、出力は 84 になります。これは、配列を次のように分割できるためです。

• [1, 15, 7] → [15, 15, 15](合計 45)
• [9] → [9](合計 9)
• [2, 5, 10] → [10, 10, 10](合計 30)

45 + 9 + 30 = 84 となり、これが達成可能な最大の合計値です。

解法アプローチ:動的計画法(DP)

この問題は動的計画法を使うことで効率的に解けます。dp[i] を「先頭から i 番目の要素までを分割したときの合計の最大値」と定義します。各位置 i について、「i 番目の要素で終わる部分配列」の長さを 1 から K まで変化させながら、その部分配列の最大値 × 長さ + 直前位置までの最適解(dp[index-1])を比較し、最大値を記録していきます。

アルゴリズムの手順

• A と同じ長さの配列 dp を作成し、0 で初期化する
• i を 0 から A の長さ - 1 まで繰り返す:
  • dp[i] = A[i] + dp[i-1](i - 1 >= 0 の場合)。それ以外は dp[i] = A[i]
  • temp := A[i] とする
  • j を 1 から K - 1 まで繰り返す:
    • i - j >= 0 の場合:
      • index := i - j
      • temp := temp と A[i - j] の大きい方
      • index - 1 >= 0 なら、dp[i] := max(dp[i], temp × (i - index + 1) + dp[index - 1])
      • そうでなければ、dp[i] := max(dp[i], temp × (i - index + 1))
• dp の最後の要素を返す

Pythonでの実装例

それでは、実際のコードを見てみましょう。

class Solution(object):
    def maxSumAfterPartitioning(self, A, K):
        dp = [0 for i in range(len(A))]
        for i in range(len(A)):
            dp[i] = A[i] + (dp[i-1] if i-1>=0 else 0)
            temp = A[i]
            for j in range(1,K):
                if i-j>=0:
                    index = i-j
                    temp = max(temp,A[i-j])
                    dp[i] = max(dp[i],temp*(i-index+1) + (dp[index-1] if index-1 >=0 else 0))
        return dp[-1]

ob = Solution()
print(ob.maxSumAfterPartitioning([1,15,7,9,2,5,10],3))

入力

[1,15,7,9,2,5,10]
3

出力

84

計算量の評価

• 時間計算量: O(N × K) — 各要素 i について、最大 K 通りの部分配列の長さを試すためです。
• 空間計算量: O(N) — DP テーブル dp を保存するために必要です。

全探索ですべての分割パターンを試すと指数時間かかりますが、DP を使うことで大幅に効率化できるのがこの手法の魅力です。