K 个不相交子数组的最大能量值

标签: 数组 动态规划 前缀和

难度: Hard

给你一个长度为 n 下标从 0 开始的整数数组 nums 和一个 正奇数 整数 k 。

x 个子数组的能量值定义为 strength = sum[1] * x - sum[2] * (x - 1) + sum[3] * (x - 2) - sum[4] * (x - 3) + ... + sum[x] * 1 ,其中 sum[i] 是第 i 个子数组的和。更正式的,能量值是满足 1 <= i <= x 的所有 i 对应的 (-1)i+1 * sum[i] * (x - i + 1) 之和。

你需要在 nums 中选择 k 个 不相交子数组 ,使得 能量值最大 。

请你返回可以得到的 最大能量值 。

注意,选出来的所有子数组  需要覆盖整个数组。

示例 1:

输入:nums = [1,2,3,-1,2], k = 3
输出:22
解释:选择 3 个子数组的最好方式是选择:nums[0..2] ,nums[3..3] 和 nums[4..4] 。能量值为 (1 + 2 + 3) * 3 - (-1) * 2 + 2 * 1 = 22 。

示例 2:

输入:nums = [12,-2,-2,-2,-2], k = 5
输出:64
解释:唯一一种选 5 个不相交子数组的方案是:nums[0..0] ,nums[1..1] ,nums[2..2] ,nums[3..3] 和 nums[4..4] 。能量值为 12 * 5 - (-2) * 4 + (-2) * 3 - (-2) * 2 + (-2) * 1 = 64 。

示例 3:

输入:nums = [-1,-2,-3], k = 1
输出:-1
解释:选择 1 个子数组的最优方案是:nums[0..0] 。能量值为 -1 。

提示:

  • 1 <= n <= 104
  • -109 <= nums[i] <= 109
  • 1 <= k <= n
  • 1 <= n * k <= 106
  • k 是奇数。

Submission

运行时间: 876 ms

内存: 17.8 MB

# https://leetcode.cn/u/l00/

class Solution:
    def maximumStrength(self, nums: List[int], k: int) -> int:
        n = len(nums)
        acc = [0]
        for num in nums:
            acc.append(acc[-1] + num)
        dps = [0] * (n + 1)
        for ck in range(k):
            cw = k - ck
            if ck & 1: cw = -cw
            # cw = (-1 if ck & 1 else 1) * (k - ck)
            lastDps = dps[ck]
            dps[ck] = curs = -inf
            # bestDps = -inf
            for i in range(ck, n - k + ck + 1):
                t = lastDps - acc[i] * cw
                if curs < t: curs = t
                lastDps = dps[i + 1]
                # if bestDps < curs + acc[i + 1] * cw: bestDps = curs + acc[i + 1] * cw
                # dps[i + 1] = bestDps
                t = curs + acc[i + 1] * cw
                dps[i + 1] = t if dps[i] < t else dps[i]
        return dps[-1]

# class Solution:
#     def maximumStrength(self, nums: List[int], k: int) -> int:
#         n = len(nums)
#         acc = 0 
#         dps = [-inf] * k
#         res = [-inf] * k
#         cws = [(k - ck) * (((ck & 1 ^ 1) << 1) -1)  for ck in range(k)]
#         for i in range(n):
#             preAcc = acc
#             acc += nums[i]
#             preRes = 0
#             bestDps = dps[0]
#             bestRes = res[0]
#             for ck in range(min(i+1, k)): # for ck in range(k):
#                 cw = cws[ck]
#                 # if preRes - preAcc * cw > bestDps: bestDps = preRes - preAcc * cw
#                 # dps[ck] = bestDps
#                 dps[ck] = max(dps[ck], preRes - preAcc * cw)
#                 preRes = res[ck]
#                 if dps[ck] + acc * cw > bestRes: bestRes = dps[ck] + acc * cw
#                 res[ck] = bestRes
#                 # dps[ck] = max(dps[ck], preRes - preAcc * cw)
#                 # preRes = res[ck]
#                 # res[ck] = max(res[ck], dps[ck] + acc * cw)
#         return res[-1]

# class Solution:
#     def maximumStrength(self, nums: List[int], k: int) -> int:
#         n = len(nums)
#         # 整理前缀和
#         acc = [0]
#         for num in nums:
#             acc.append(acc[-1] + num)
#         # 动态规划每一层选出下标 i 范围内划分 ck 块的最优结果
#         preDps = [0] * (n + 1)
#         dps = [0] * (n + 1)
#         for ck in range(k):
#             cw = k - ck
#             if ck & 1: cw = -cw
#             dps[ck] = -inf
#             curs = -inf
#             for i in range(ck, n):
#                 curs = max(curs, preDps[i] - acc[i] * cw)
#                 dps[i + 1] = max(dps[i], curs + acc[i + 1] * cw)
#             dps, preDps = preDps, dps
#         return preDps[-1]

Explain

本题解采用动态规划的方式来求解。首先,计算数组 `nums` 的累加和 `acc`,这是为了快速计算任意子数组的和。然后,定义动态规划数组 `dps`,其中 `dps[i]` 表示在数组前 `i` 个元素中选取 `ck` 个子数组能够获得的最大能量值。对于每一个可能的 `k`,计算对应的权重 `cw`,这个权重根据 `k` 的值和当前子数组的索引 `ck` 而有所不同,具体为 `(-1)^ck * (k - ck)`。接着,遍历数组使用滚动数组方式更新 `dps`,每次考虑加入新的子数组时,更新 `dps[i]` 为包括新子数组的最大能量值。这个过程中,我们需要维护当前的最大值 `curs` 以便快速更新 `dps`。最终 `dps[-1]` 将包含在整个数组中选取 `k` 个子数组的最大能量值。

时间复杂度: O(n^2)

空间复杂度: O(n)

# Class definition for the solution to the problem

class Solution:
    def maximumStrength(self, nums: List[int], k: int) -> int:
        n = len(nums)
        # Create prefix sum array
        acc = [0]
        for num in nums:
            acc.append(acc[-1] + num)
        # Initialize the DP array
        dps = [0] * (n + 1)
        # Iterate through possible number of subarrays
        for ck in range(k):
            # Calculate the weight for current number of subarrays
            cw = k - ck
            if ck & 1: cw = -cw
            lastDps = dps[ck]
            dps[ck] = curs = -inf
            # Update DP values for the current subarray count
            for i in range(ck, n - k + ck + 1):
                t = lastDps - acc[i] * cw
                if curs < t: curs = t
                lastDps = dps[i + 1]
                t = curs + acc[i + 1] * cw
                dps[i + 1] = t if dps[i] < t else dps[i]
        # Return the maximum strength from the last DP value
        return dps[-1]

Explore

动态规划是解决此类问题的有效方法,因为它能够将问题分解成较小的子问题,并通过解决这些子问题来构建最终解决方案。特别是在涉及到多个选择和决策步骤的问题中,如在数组中选择不相交子数组以最大化某种指标(本题中的能量值)时,动态规划可以有效地避免重复计算并储存中间结果,从而降低问题的复杂度。此外,动态规划通过迭代更新解决方案,能够确保每一步都是基于前一步的最优解,这使得最终得到的是全局最优解。

在算法的实现中,`dps[i]` 记录的是从数组的开始到第 `i` 个元素,最多取 `ck` 个子数组时的最大能量值。通过逐步增长子数组的数量并更新 `dps[i]`,在每一步中,算法都会检查前一个最优结果和当前可能加入的子数组的组合。这里的关键在于,每次更新 `dps[i]` 时,都是基于在第 `i` 个元素之前的数组段的最优解进行更新。因此,新增的子数组始终是从前一个子数组结束之后的位置开始,这自然地保证了子数组之间不会有重叠。

权重 `cw` 的计算方式是为了根据子数组的数量调整每个子数组对总能量值的贡献。这里的计算公式 `(-1)^ck * (k - ck)` 是基于子数组数量来调整权重的正负和大小。当 `ck` (当前子数组的数量)是奇数时,权重为负,这可能代表在某些情境下需要减少由于选择更多子数组带来的额外成本或惩罚。这种权重调整方式确保算法在增加子数组数量时能够平衡每个子数组的贡献,从而在满足题目要求的同时最大化总能量值。在整个解决方案的上下文中,这样的权重调整帮助算法在不同的选择中找到最优解。