Pythonで石の山を1つにまとめる最小コストを求めるプログラム(区間DP)

問題の概要

一列に並んだ N 個の石の山があり、i 番目の山には stones[i] 個の石が入っています。1回の操作では「連続する K 個の山」を1つの山にまとめることができ、そのときのコストは K 個の山に含まれる石の総数に等しくなります。すべての山を1つにまとめる際の最小コストを求めてください。ただし、まとめ方が存在しない場合は -1 を返します。

具体例

nums = [3, 2, 4, 1]、K = 2 の場合、出力は 20 になります。

• 初期状態:[3, 2, 4, 1]
• [3, 2] をマージ(コスト 5)→ [5, 4, 1]
• [4, 1] をマージ(コスト 5)→ [5, 5]
• [5, 5] をマージ(コスト 10)→ [10]

合計コストは 5 + 5 + 10 = 20 となり、これが最小値です。

解法のポイント

まず押さえておきたい性質として、1回のマージを行うたびに山の数は必ず K − 1 個減ります。したがって、n 個の山を最終的に1つにするためには、(n − 1) が (K − 1) の倍数でなければなりません。この条件を満たさない入力に対しては、答えとして -1 を返します。

区間DPによるアプローチ

dp[i][j] を「区間 [i, j] 内の山を、可能な限り少ない数の山にまとめるときの最小コスト」と定義します。(j − i) mod (K − 1) が 0 になる場合、その区間はちょうど1つの山にまとめられるため、区間全体の石の総数(累積和の差分)をコストに加算します。

アルゴリズムの手順は以下の通りです。

1. n を nums のサイズとします。
2. (n − 1) mod (K − 1) が 0 でない場合は -1 を返します。
3. dp を n × n の2次元配列として用意し、0 で初期化します。
4. sums をサイズ n + 1 の累積和配列として用意し、0 で初期化します。
5. i が 1 から n までの範囲で、sums[i] = sums[i−1] + nums[i−1] を計算します。
6. length を K から n まで増やしながら、以下を繰り返します。
   • i を 0 から n − length まで動かし、j = i + length − 1 とします。
   • dp[i][j] を無限大で初期化します。
   • t を i から j − 1 まで K − 1 ずつ増やしながら、dp[i][j] = min(dp[i][j], dp[i][t] + dp[t+1][j]) で更新します。
   • (j − i) mod (K − 1) が 0 の場合、dp[i][j] に sums[j+1] − sums[i](区間の石の総数)を加算します。
7. 最後に dp[0][n−1] を返します。

Pythonでの実装例

以下が実際の Python コードです。

def solve(nums, K):
    n = len(nums)
    if (n - 1) % (K - 1) != 0:
        return -1
    dp = [[0] * n for _ in range(n)]
    sums = [0] * (n + 1)
    for i in range(1, n + 1):
        sums[i] = sums[i - 1] + nums[i - 1]
    for length in range(K, n + 1):
        for i in range(n - length + 1):
            j = i + length - 1
            dp[i][j] = float('inf')
            for t in range(i, j, K - 1):
                dp[i][j] = min(dp[i][j], dp[i][t] + dp[t + 1][j])
            if (j - i) % (K - 1) == 0:
                dp[i][j] += sums[j + 1] - sums[i]
    return dp[0][n - 1]

nums = [3, 2, 4, 1]
K = 2
print(solve(nums, K))

入力

[3,2,4,1], 2

出力

20

計算量

時間計算量は O(n³ / K)、空間計算量は O(n²) です。内側の t に関するループが K − 1 刻みで進むため、素朴な区間DPよりも効率的に動作します。