合并石头的最低成本 -- 三维动规数组

题目描述

题目链接：1000. 合并石头的最低成本

有 N 堆石头排成一排，第 i 堆中有 stones[i] 块石头。
每次移动（move）需要将连续的 K 堆石头合并为一堆，而这个移动的成本为这 K 堆石头的总数。
找出把所有石头合并成一堆的最低成本。如果不可能，返回 -1 。

代码实现

class Solution:
    def mergeStones(self, arr: List[int], k: int) -> int:
        n = len(arr)
        if (n-1) % (k-1) != 0:
            return -1

        pre_sum = [0]*(n+1)
        for i in range(n):
            pre_sum[i+1] += pre_sum[i] + arr[i]
        dp = [[[0]*(k+1) for _ in range(n)] for _ in range(n)]
        return self.process3(0, n-1, 1, arr, k, pre_sum, dp)

    def process3(self, L, R, part, arr, k, pre_sum, dp):

        """
        arr[L..R] 一定要弄出part份，返回最低代价
        :param L:
        :param R:
        :param part:
        :param arr:
        :param k:
        :param pre_sum:
        :return:
        """
        if dp[L][R][part] != 0:
            return dp[L][R][part]

        if L == R:
            return 0 if part == 1 else -1
        # L ... R 不只一个数
        if part == 1:
            next_ = self.process3(L, R, k, arr, k, pre_sum, dp)
            if next_ == -1:
                return -1
            else:
                dp[L][R][part] = next_ + pre_sum[R+1] - pre_sum[L]
                return next_ + pre_sum[R+1] - pre_sum[L]
        else:
            ans = float('inf')
            for mid in range(L, R, k-1):
                # L..mid(一份); mid+1...R(part - 1)
                next1 = self.process3(L, mid, 1, arr, k, pre_sum, dp)
                next2 = self.process3(mid+1, R, part-1, arr, k, pre_sum, dp)
                if next1 != -1 and next2 != -1:
                    ans = min(ans, next1+next2)
            dp[L][R][part] = ans
            return ans