Problem: 3013. 将数组分成最小总代价的子数组 II
思路
第一个数是明确的,即nums[0]。
我们接下来要求的某个连续区间且长度为dist + 1
,它的前k - 1
个最小的数的和最小,我们的答案就是`nums[0] + min(arr[i ~ i + dist + 1]),min是求arr的前k小的和。
解题方法
这里我们可以去回顾
lc.480
这道题,官方题解写的非常好的类模板,我写了一些注释,每个模块和函数都显而易见,思路很清晰。不过这道题仅仅是求一个滑动窗口中的中位数,而我们这题是求前k个小数的和。
因此我们这里需要多维护一个值prefix
为答案,就可以套用模板了。
这模板虽然看着有些长,但是自己跟着去理解然后敲一遍,还是很容易默写的。
以后类似的题,都可以用这套模板,大部分只需要稍微改变一下makeBalance、insert、erase函数。
模板如下:
class DualHeap {
private:
// 下堆是大根堆
priority_queue<int> down;
// 上堆是小根堆
priority_queue<int, vector<int>, greater<int>> up;
// 延迟删除堆中元素,只有当其在堆顶的时候才进行删除
// [k, v] 分别代表[要删除的数值、出现的次数]
unordered_map<int, int> delayed;
// 窗口大小(题目要求)
int k;
// 这里十分重要,调整平衡的时候不是比较虚拟大小(down.size),而是比较真实大小(***重点***)
// 记录上下堆的实际大小
int downSize, upSize;
public:
DualHeap(int _k): k(_k), downSize(0), upSize(0) {}
private:
// 延迟删除(修剪)
// 如果堆顶元素是要删除的元素,则一直删除,直到堆顶元素在滑窗内
void prune(auto &heap) {
while (heap.size()) {
int x = heap.top();
if (delayed.count(x)) {
if (-- delayed[x] == 0)
delayed.erase(x);
heap.pop();
}
else break;
}
}
// 维持两个堆的大小
/** 维持条件:
1.下堆至多大于上堆1个,至少与上堆个数相等
2.min(上堆) >= max(下堆)
**/
void makeBalance() {
if (downSize > upSize + 1) {
up.push(down.top());
down.pop();
-- downSize, ++ upSize;
prune(down);
}
else if (downSize < upSize) {
down.push(up.top());
up.pop();
++ downSize, -- upSize;
prune(up);
}
}
public:
// 插入
void insert(int num) {
if (down.empty() || num <= down.top()) down.push(num), ++ downSize;
else up.push(num), ++ upSize;
makeBalance();
}
// 删除
void erase(int num) {
++ delayed[num];
if (num <= down.top()) {
-- downSize;
if (num == down.top())
prune(down);
}
else {
-- upSize;
if (num == up.top())
prune(up);
}
makeBalance();
}
// 求中位数(题目要求)
double getMedian() {
return k & 1 ? down.top() : ((double)down.top() + up.top()) / 2;
}
};
复杂度
时间复杂度:
总的时间复杂度为 O ( n l o g n ) O(nlogn) O(nlogn),外层遍历 O ( n ) O(n) O(n),内层堆的每一次push为 O ( l o g n ) O(logn) O(logn)
空间复杂度:
辅助空间用到优先队列为 O ( n ) O(n) O(n)
Code
typedef long long LL;
class DualHeap {
private:
priority_queue<int> down;
priority_queue<int, vector<int>, greater<int>> up;
int dist, k;
LL prefix;
unordered_map<int, int> delayed;
int downSize, upSize;
public:
DualHeap(int _dist, int _k): dist(_dist), k(_k),
downSize(0), upSize(0), prefix(0LL) {}
private:
// 与模板一致
void prune(auto &heap) {
while (heap.size()) {
int x = heap.top();
if (delayed.count(x)) {
if (-- delayed[x] == 0) delayed.erase(x);
heap.pop();
}
else break;
}
}
// 这里有些不同
void makeBalance() {
// 如果down中的真实元素个数大于k-1,取下往上插入
if (downSize > k - 1) {
up.push(down.top());
// 调整平衡的时候,取下往上插,这里要减去
prefix -= down.top();
down.pop();
downSize --, upSize ++;
prune(down);
}
// 这里记得判断up不为空
else if (downSize < k - 1 && up.size()) {
down.push(up.top());
// 同样的,调整平衡的时候,取上往下插,这里要加上
prefix += up.top();
up.pop();
downSize ++, upSize --;
prune(up);
}
}
public:
// 插入的时候,一定是真实元素,要加上
void insert(int num) {
if (down.empty() || num <= down.top())
down.push(num), ++ downSize, prefix += num;
else up.push(num), upSize ++;
makeBalance();
}
// 删除的时候,一定是删除真实元素,如果num小于down.top()说明它一定在下堆,直接减
void erase(int num) {
delayed[num] ++;
if (num <= down.top()) {
downSize --, prefix -= num;
if (num == down.top())
prune(down);
} else {
upSize --;
if (num == up.top())
prune(up);
}
makeBalance();
}
LL getPrefixSumOfk() {
return prefix;
}
};
class Solution {
public:
long long minimumCost(vector<int>& nums, int k, int dist) {
int n = nums.size();
DualHeap dh(dist, k);
for (int i = 1; i <= dist + 1; i ++) dh.insert(nums[i]);
// 这别忘记了
LL res = nums[0] + dh.getPrefixSumOfk();
for (int i = dist + 2; i < n; i ++) {
dh.insert(nums[i]);
dh.erase(nums[i - dist - 1]);
res = min(res, dh.getPrefixSumOfk() + nums[0]);
}
return res;
}
};