题目描述
给你一个长度为 n
下标从 0 开始的正整数数组 nums
。
同时给你一个长度为 m
的二维操作数组 queries
,其中 queries[i] = [index_i, k_i]
。
一开始,数组中的所有元素都 未标记。
你需要依次对数组执行 m
次操作,第 i
次操作中,你需要执行:
- 如果下标
index_i
对应的元素还没标记,那么标记这个元素。 - 然后标记
k_i
个数组中还没有标记的 最小 元素。如果有元素的值相等,那么优先标记它们中下标较小的。如果少于k_i
个未标记元素存在,那么将它们全部标记。
请你返回一个长度为 m
的数组 answer
,其中 answer[i]
是第 i
次操作后数组中还没标记元素的 和。
样例
输入: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。
输入: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 <= 10^5
1 <= n <= 10^5
queries[i].length == 2
0 <= index_i, k_i <= n - 1
算法
(堆) $O(n \log n + m)$
- 维护一个小根堆,存储元素值和其下标,以及一个标记数组,使用标记数组进行堆的延迟删除。
- 每次询问,标记
index
,然后从堆中弹出最小未标记的k
个元素,如果当前弹出的元素已经被标记了,则当前这个元素不计入本次的k
个元素。
时间复杂度
- 构造堆的时间复杂度为 $O(n \log n)$。
- 所有询问累加最多弹出 $n$ 个元素,总时间复杂度为 $O(n \log n)$。
- 故总时间复杂度为 $O(n \log n + m)$。
空间复杂度
- 需要 $O(n + m)$ 的额外空间存储堆,标记数组和答案。
C++ 代码
#define LL long long
class Solution {
public:
vector<LL> unmarkedSumArray(vector<int>& nums, vector<vector<int>>& queries) {
const int n = nums.size(), m = queries.size();
vector<LL> ans(m);
LL tot = 0;
priority_queue<pair<int, int>> heap;
for (int i = 0; i < n; i++) {
tot += nums[i];
heap.push(make_pair(-nums[i], -i));
}
vector<bool> seen(n, false);
for (int i = 0; i < m; i++) {
int idx = queries[i][0], k = queries[i][1];
if (!seen[idx]) {
tot -= nums[idx];
seen[idx] = true;
}
while (!heap.empty() && k > 0) {
auto u = heap.top();
heap.pop();
if (!seen[-u.second]) {
seen[-u.second] = true;
tot += u.first;
--k;
}
}
ans[i] = tot;
}
return ans;
}
};