将数组分成最小总代价的子数组 II

难度: Hard

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

一个数组的代价是数组中的 第一个 元素。比方说，[1,2,3] 的代价为 1 ，[3,4,1] 的代价为 3 。

你需要将 nums 分割成 k 个 连续且互不相交 的子数组，满足第二个子数组与第 k 个子数组中第一个元素的下标距离 不超过 dist 。换句话说，如果你将 nums 分割成子数组 nums[0..(i₁ - 1)], nums[i₁..(i₂ - 1)], ..., nums[i_k-1..(n - 1)] ，那么它需要满足 i_k-1 - i₁ <= dist 。

请你返回这些子数组的最小总代价。

示例 1：

输入：nums = [1,3,2,6,4,2], k = 3, dist = 3
输出：5
解释：将数组分割成 3 个子数组的最优方案是：[1,3] ，[2,6,4] 和 [2] 。这是一个合法分割，因为 i_k-1 - i₁ 等于 5 - 2 = 3 ，等于 dist 。总代价为 nums[0] + nums[2] + nums[5] ，也就是 1 + 2 + 2 = 5 。
5 是分割成 3 个子数组的最小总代价。

示例 2：

输入：nums = [10,1,2,2,2,1], k = 4, dist = 3
输出：15
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[1] ，[2] 和 [2,2,1] 。这是一个合法分割，因为 i_k-1 - i₁ 等于 3 - 1 = 2 ，小于 dist 。总代价为 nums[0] + nums[1] + nums[2] + nums[3] ，也就是 10 + 1 + 2 + 2 = 15 。
分割 [10] ，[1] ，[2,2,2] 和 [1] 不是一个合法分割，因为 i_k-1 和 i₁ 的差为 5 - 1 = 4 ，大于 dist 。
15 是分割成 4 个子数组的最小总代价。

示例 3：

输入：nums = [10,8,18,9], k = 3, dist = 1
输出：36
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[8] 和 [18,9] 。这是一个合法分割，因为 i_k-1 - i₁ 等于 2 - 1 = 1 ，等于 dist 。总代价为 nums[0] + nums[1] + nums[2] ，也就是 10 + 8 + 18 = 36 。
分割 [10] ，[8,18] 和 [9] 不是一个合法分割，因为 i_k-1 和 i₁ 的差为 3 - 1 = 2 ，大于 dist 。
36 是分割成 3 个子数组的最小总代价。

提示：

3 <= n <= 10⁵
1 <= nums[i] <= 10⁹
3 <= k <= n
k - 2 <= dist <= n - 2

Submission

运行时间: 370 ms

内存: 30.8 MB

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

class Solution:
    def minimumCost(self, nums: List[int], k: int, dist: int) -> int:
        n = len(nums)

        removesBook = defaultdict(int)

        pickedHeap = [-num for num in nums[1:k]]
        heapify(pickedHeap)

        unpickedHeap = []
        bestRes = curRes = sum(nums[:k])
        
        for right in range(k, n):
            left = right - dist

            # 判断添加元素入已选堆还是备选堆
            newNum = nums[right]
            if newNum < -pickedHeap[0]:
                maxNum = heappushpop(pickedHeap, -newNum)
                heappush(unpickedHeap, -maxNum)
                curRes += newNum + maxNum
            else:
                heappush(unpickedHeap, newNum)

            # 旧窗口左侧，即 left - 1 数值需要登记扣除
            if left > 1:
                outNum = nums[left - 1]
                removesBook[outNum] += 1
                if outNum <= -pickedHeap[0]:
                    while unpickedHeap and removesBook[unpickedHeap[0]]:
                        removesBook[heappop(unpickedHeap)] -= 1
                    if outNum < -pickedHeap[0] or removesBook[outNum]:
                        pickNum = heappop(unpickedHeap)
                        heappush(pickedHeap, -pickNum)
                        curRes += pickNum - outNum
                # 整理堆头
                while removesBook[-pickedHeap[0]]:
                    removesBook[-heappop(pickedHeap)] -= 1

            # 总额更新
            if curRes < bestRes: bestRes = curRes
                
        return bestRes

Explain

该题解使用了两个堆（最大堆和最小堆）以及一个哈希表来维护和更新窗口中的元素，其中窗口大小由 dist 控制。窗口代表了可以作为子数组起始点的范围。解题思路的核心在于动态地维护k个子数组的最小总代价，通过在窗口中添加新元素，并从窗口移除旧元素。最大堆（pickedHeap）用于维护已选择的元素（负值表示，因为Python的heapq实现的是最小堆），而最小堆（unpickedHeap）用于维护未选择的元素。当一个新元素大于最大堆的堆顶时，它被加入到最小堆；反之，则替换最大堆的堆顶元素，并将旧的堆顶元素移到最小堆。这样可以保证当前窗口内的k个元素中，最小的元素总是位于最大堆的堆顶。哈希表（removesBook）用于记录已经从窗口中移除的元素的计数，以便在更新堆时正确处理这些元素。

时间复杂度: O(n log k)

空间复杂度: O(n)

# Solution to split array into k subarrays with minimum total cost and constraints

from heapq import heappush, heappop, heapify
from collections import defaultdict
from typing import List

class Solution:
    def minimumCost(self, nums: List[int], k: int, dist: int) -> int:
        n = len(nums)
        # Initializes heaps and bookkeeping dictionary
        removesBook = defaultdict(int)
        pickedHeap = [-num for num in nums[1:k]] # max-heap (using min-heap with negative values)
        heapify(pickedHeap)
        unpickedHeap = [] # min-heap
        bestRes = curRes = sum(nums[:k]) # initialize with the first k elements cost
        
        for right in range(k, n):
            left = right - dist
            # Decides whether to add to the max-heap or min-heap
            newNum = nums[right]
            if newNum < -pickedHeap[0]:
                maxNum = heappushpop(pickedHeap, -newNum)
                heappush(unpickedHeap, -maxNum)
                curRes += newNum + maxNum
            else:
                heappush(unpickedHeap, newNum)
            # Manage elements that move out of the window
            if left > 1:
                outNum = nums[left - 1]
                removesBook[outNum] += 1
                if outNum <= -pickedHeap[0]:
                    while unpickedHeap and removesBook[unpickedHeap[0]]:
                        removesBook[heappop(unpickedHeap)] -= 1
                    if outNum < -pickedHeap[0] or removesBook[outNum]:
                        pickNum = heappop(unpickedHeap)
                        heappush(pickedHeap, -pickNum)
                        curRes += pickNum - outNum
                while removesBook[-pickedHeap[0]]:
                    removesBook[-heappop(pickedHeap)] -= 1
            # Update the minimum result
            if curRes < bestRes: bestRes = curRes
        
        return bestRes

Explore

在算法中使用两个堆可以有效地维护当前窗口内元素的最小总代价。最大堆（使用负值实现）用于追踪窗口中的最小元素，因为这些元素直接影响总代价。最小堆则用来存储未被选为最小元素的其他元素。通过这种方式，每当窗口移动或元素更新时，可以快速地从最大堆中获取最小元素，同时将新的较小元素加入最大堆，并将原最小元素移至最小堆，从而保证在动态变化的窗口中总是能够迅速地计算和更新最小总代价。

使用哈希表记录已从窗口中移除的元素的计数是为了正确管理和更新堆中的元素。由于堆结构本身不支持随机访问，当窗口滑动导致某些元素不再属于当前窗口时，这些元素可能仍在堆中。通过哈希表记录这些元素的计数，我们可以在执行堆操作（如弹出元素）时检查顶部元素是否已被移出窗口，如果是，则将其弹出并调整堆，直到堆顶是有效的窗口元素。这样确保了堆中的数据始终反映当前窗口的状态，从而正确计算最小总代价。

堆操作在算法中确保了每次窗口更新后，最大堆始终包含当前窗口中的最小元素，从而保持了k个子数组的最小总代价。具体来说，每当新元素加入窗口时，会与最大堆的堆顶（当前最小元素）比较，根据比较结果决定是替换堆顶还是加入最小堆。当窗口移动导致元素退出时，通过哈希表和堆操作来确保移除的是正确的元素，并从最小堆中补充元素到最大堆，保持最大堆的大小和堆顶元素的正确性。通过这些操作，算法可以持续调整并保持最小总代价。