执行操作标记数组中的元素

标签: 数组 哈希表 排序 模拟 堆（优先队列）

难度: Medium

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

同时给你一个长度为 m 的二维操作数组 queries ，其中 queries[i] = [indexi, ki] 。

一开始，数组中的所有元素都 未标记 。

你需要依次对数组执行 m 次操作，第 i 次操作中，你需要执行：

  • 如果下标 indexi 对应的元素还没标记，那么标记这个元素。
  • 然后标记 ki 个数组中还没有标记的 最小 元素。如果有元素的值相等，那么优先标记它们中下标较小的。如果少于 ki 个未标记元素存在，那么将它们全部标记。

请你返回一个长度为 m 的数组 answer ，其中 answer[i]是第 i 次操作后数组中还没标记元素的 和 。

示例 1：

输入：nums = [1,2,2,1,2,3,1], queries = [[1,2],[3,3],[4,2]]

输出：[8,3,0]

解释：

我们依次对数组做以下操作：

  • 标记下标为 1 的元素，同时标记 2 个未标记的最小元素。标记完后数组为 nums = [1,2,2,1,2,3,1] 。未标记元素的和为 2 + 2 + 3 + 1 = 8 。
  • 标记下标为 3 的元素，由于它已经被标记过了，所以我们忽略这次标记，同时标记最靠前的 3 个未标记的最小元素。标记完后数组为 nums = [1,2,2,1,2,3,1] 。未标记元素的和为 3 。
  • 标记下标为 4 的元素，由于它已经被标记过了，所以我们忽略这次标记，同时标记最靠前的 2 个未标记的最小元素。标记完后数组为 nums = [1,2,2,1,2,3,1] 。未标记元素的和为 0 。

示例 2：

输入：nums = [1,4,2,3], queries = [[0,1]]

输出：[7]

解释：我们执行一次操作，将下标为 0 处的元素标记，并且标记最靠前的 1 个未标记的最小元素。标记完后数组为 nums = [1,4,2,3] 。未标记元素的和为 4 + 3 = 7 。

提示：

  • n == nums.length
  • m == queries.length
  • 1 <= m <= n <= 105
  • 1 <= nums[i] <= 105
  • queries[i].length == 2
  • 0 <= indexi, ki <= n - 1

Submission

运行时间: 257 ms

内存: 38.0 MB

class Solution:
    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        n = len(nums)
        s = sum(nums)
        ids = sorted(range(n), key=lambda i: nums[i])  # 稳定排序
        ans = []
        j = 0
        for i, k in queries:
            s -= nums[i]
            nums[i] = 0  # 标记
            while j < n and k:
                i = ids[j]
                if nums[i]:  # 没有被标记
                    s -= nums[i]
                    nums[i] = 0
                    k -= 1
                j += 1
            ans.append(s)
        return ans

Explain

题解首先通过sum函数计算整个nums数组的总和s，然后创建一个索引数组ids，根据nums中元素的值进行稳定排序。这样，ids数组的每个位置保存的是原nums数组中元素从小到大排序的索引。接着，题解定义一个变量j来跟踪未标记的最小元素的位置。对于每个查询，题解首先从总和s中减去指定索引的元素值（即标记该元素），然后尝试标记k个未标记的最小元素，通过索引数组ids来找到这些元素。如果元素已经被标记（即值为0），则忽略。每次操作后，将当前未标记元素的和s添加到结果列表ans中。

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

空间复杂度: O(n)

class Solution:
    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        n = len(nums)  # 获取数组长度
        s = sum(nums)  # 计算数组总和
        ids = sorted(range(n), key=lambda i: nums[i])  # 根据元素值稳定排序索引
        ans = []  # 初始化结果列表
        j = 0  # 初始化未标记的最小元素的位置
        for i, k in queries:
            s -= nums[i]  # 从总和中减去当前操作的元素
            nums[i] = 0  # 标记当前元素
            while j < n and k:  # 继续标记k个未标记的最小元素
                i = ids[j]  # 获取未标记的最小元素的索引
                if nums[i]:  # 如果当前元素未被标记
                    s -= nums[i]  # 从总和中减去该元素
                    nums[i] = 0  # 标记该元素
                    k -= 1  # 减少需要标记的元素数量
                j += 1  # 移动到下一个最小元素
            ans.append(s)  # 记录当前未标记元素的总和
        return ans  # 返回结果

Explore

在题解中，`ids`数组是通过`sorted(range(n), key=lambda i: nums[i])`产生的。这里，`range(n)`生成一个与`nums`数组索引对应的序列，`key=lambda i: nums[i]`指定排序依据为`nums`数组中索引`i`对应的值。Python的`sorted`函数默认实现的是稳定排序，即如果两个元素具有相同的排序键（在本例中为`nums[i]`），它们在排序后的数组中的相对位置与在原数组中的相对位置相同。因此，`ids`数组确保了元素的稳定排序。

题解中处理这种情况的逻辑位于`if nums[i]:`条件判断内。当进行到某次查询操作时，如果发现`nums[i]`的值为0，说明该元素已经被之前的操作标记过，因此，程序将不再从`s`中减去`nums[i]`的值（因为已经是0），直接继续检查下一个元素。这确保了只有未标记的元素才会影响`s`的值，并且准确反映了每次操作后未标记元素的总和。

题解中的`while j < n and k`循环确实处理了所有元素可能已经标记的情况。循环中，`j`遍历`ids`数组（按元素值排序的索引数组），并检查`nums[i]`是否已被标记（`nums[i]`不为0）。如果在循环执行过程中`j`达到`n`（即检查完所有元素）但`k`仍然大于0，循环将自然结束，因为没有更多元素可供标记。此时，`s`已经是剩余未标记元素的总和，直接添加到`ans`列表即可，不需要额外操作。