3013. 將數組分成最小總代價的子數組 II

困難 數組 哈希表 滑動窗口

前置題目:3010. 將數組分成最小總代價的子數組 I

思路

把陣列切成 kk 份,第一個陣列的代價是 nums[0]
現在還要挑其他 k1k-1 個位置作為其他子陣列的起始點。
額外的條件限制是第二個子陣列的起點跟最後一個陣列的起點距離不能超過dist

這也就代表,現在要在陣列中找一個長度為 dist+1 的窗口,
在窗口裡面選出 k-1 個元素,讓它們的合最小。
因此在移動窗口去加入與移除數字時,要時刻保持最小的 k-1 個元素
可以用有序集合或優先隊列去做到這點。


建立兩個有序集合,smalllarge

  • 窗口加入數字時,先將數字加到 small 中,確保裝小數字的 small 包含 k-1 個元素,如果 small 滿了,就丟到 large 裡面
  • 縮減窗口時, 先找出數字在 small 還是 large,並丟掉,如果 small 元素過少,就從 large 取最小數字補充到 small 中。

在過程當中,紀錄 small 的數值總和。


假如使用優先隊列,把數值在 smalllarge 搬來搬去的時候,
需要紀錄當前最新的狀態,避免過時的資訊影響操作,這個技巧叫做懶刪除堆。

就這題而言,優先隊列雖然比較快,但寫得很痛苦,建議用第一種方法。

程式碼

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;
    }
};

複雜度分析

  • 時間複雜度:O(nlogn)O(n\log{n})
  • 空間複雜度:O(n)O(n)

顯示設定

背景線條
顯示背景網格線條
懸停發光
滑鼠懸停時顯示霓虹效果
聚光燈
跟隨滑鼠的聚光燈效果
背景透明度
開啟透明玻璃效果
主題顏色
自訂主要顏色