对顶堆+延迟删除(lc.3013)

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(nlogn)$，外层遍历$O(n)$，内层堆的每一次push为$O(logn)$

空间复杂度:

辅助空间用到优先队列为$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;
    }
};