思路
把陣列切成 份,第一個陣列的代價是 nums[0],
現在還要挑其他 個位置作為其他子陣列的起始點。
額外的條件限制是第二個子陣列的起點跟最後一個陣列的起點距離不能超過dist。
這也就代表,現在要在陣列中找一個長度為 dist+1 的窗口,
在窗口裡面選出 k-1 個元素,讓它們的合最小。
因此在移動窗口去加入與移除數字時,要時刻保持最小的 k-1 個元素
可以用有序集合或優先隊列去做到這點。
建立兩個有序集合,small跟large,
- 窗口加入數字時,先將數字加到
small中,確保裝小數字的small包含k-1個元素,如果small滿了,就丟到large裡面 - 縮減窗口時, 先找出數字在
small還是large,並丟掉,如果small元素過少,就從large取最小數字補充到small中。
在過程當中,紀錄 small 的數值總和。
假如使用優先隊列,把數值在 small 跟 large 搬來搬去的時候,
需要紀錄當前最新的狀態,避免過時的資訊影響操作,這個技巧叫做懶刪除堆。
就這題而言,優先隊列雖然比較快,但寫得很痛苦,建議用第一種方法。
程式碼
1. 有序集合
class Solution {
private:
void smallToLarge(multiset<int>& small, multiset<int>& large, long long& smallSum) {
int val = *small.rbegin();
smallSum -= val;
small.erase(prev(small.end())); // multiset 如果移除 val, 相同數值會全部被移除,所以要用迭代器
large.insert(val);
}
void largeToSmall(multiset<int>& small, multiset<int>& large, long long& smallSum) {
if(large.empty()) return;
int val = *large.begin();
smallSum += val;
large.erase(large.begin());
small.insert(val);
}
public:
long long minimumCost(vector<int>& nums, int k, int dist) {
int n = nums.size();
long long res = LLONG_MAX;
long long smallSum = 0;
multiset<int> small, large;
for(int r = 1; r < n; r++) { // 忽略 nums[0]
// 1. 加入新元素
small.insert(nums[r]);
smallSum += nums[r];
// 如果數字太多,需要移動到 large 當中
if(small.size() == k) {
smallToLarge(small, large, smallSum);
}
// 窗口大小還不到 dist + 1,不需要縮減
if(r < dist + 1) {
continue;
}
// 此時窗口大小為 dist + 1,最小的 k-1 個數字總和列入參考答案
res = min(res, smallSum);
// 2. 移除舊元素
int l = r - dist;
auto it = small.find(nums[l]);
if(it != small.end()) { // 刪除的數字在 small
small.erase(it);
smallSum -= nums[l];
largeToSmall(small, large, smallSum); // 補充一個數字回來
}
else { // 刪除的數字在 large
large.erase(large.find(nums[l]));
}
}
return 1LL * nums[0] + res;
}
};
2. 優先隊列 + 懶刪除堆
class Solution {
public:
long long minimumCost(vector<int>& nums, int k, int dist) {
int n = nums.size();
if (k == 1) return nums[0];
int need = k - 1;
using pii = pair<int,int>; // {val, idx},解決重複數值的刪除問題
priority_queue<pii> small;
priority_queue<pii, vector<pii>, greater<pii>> large;
vector<char> dead(n, 0); // 懶刪除標記: 紀錄哪些 index 已經滑出窗口
long long smallSum = 0; // 當前 small 堆內有效元素的總和
long long res = LLONG_MAX;
int smallSize = 0, largeSize = 0; // 追蹤有效元素數量 (因為堆的 .size() 包含已失效元素,不能用)
// 清理已經「失效」(不在窗口內) 的元素
auto clearSmall = [&](){
while(!small.empty() && dead[small.top().second]) {
small.pop();
}
};
auto clearLarge = [&](){
while(!large.empty() && dead[large.top().second]) {
large.pop();
}
};
// 維持數量平衡:確保 small 剛好有 need 個最小元素 (對應 smallToLarge / largeToSmall)
auto rebalance = [&](){
clearSmall();
clearLarge();
// 如果 small 太多了,把最大的移到 large
while(smallSize > need) {
auto [val, idx] = small.top();
small.pop();
smallSum -= val;
smallSize--;
large.push({val, idx});
largeSize++;
clearSmall(); // 每次彈出後都要清理堆頂
}
// 如果 small 不夠,從 large 補最小的過來
while(smallSize < need && largeSize > 0) {
clearLarge();
if (large.empty()) break;
auto [val, idx] = large.top();
large.pop();
smallSum += val;
smallSize++;
largeSize--;
small.push({val, idx});
clearLarge();
}
};
for(int r = 1; r < n; r++){
// 1. 進:新增 nums[r]
clearSmall();
// 決定該放入哪個堆:如果比 small 最大的還小,就進 small
if(smallSize == 0 || nums[r] <= small.top().first){
small.push({nums[r], r});
smallSum += nums[r];
smallSize++;
} else {
large.push({nums[r], r});
largeSize++;
}
rebalance();
// 窗口還沒達到 dist+1 之前不結算
if(r < dist + 1) continue;
// 2. 結算:當前窗口為 [r-dist, r],記錄 smallSum 為候選答案
res = min(res, smallSum);
// 3. 出:移除滑出窗口的 nums[l]
int l = r - dist;
dead[l] = 1; // 標記 index 死掉了
// 判斷被移除的數字在哪個堆:這會影響 smallSum 和 smallSize
clearSmall();
// 如果被移除的索引對應的值 <= small 堆頂,我們認定它在 small 中
if(smallSize > 0 && nums[l] <= small.top().first) {
// 注意:這裡直接減去 nums[l],懶刪除會在堆頂遇到它時才 pop
smallSum -= nums[l];
smallSize--;
}
else {
largeSize--;
}
// 移除後可能破壞平衡,再次調整
rebalance();
}
return 1LL * nums[0] + res;
}
};
複雜度分析
- 時間複雜度:
- 空間複雜度: