從2023濟南K學習滑動視窗中位數問題

加固文明幻景發表於2024-10-03

板子

  • 對頂堆
template<class T>
struct DualHeap {
 
    Heap<T, std::greater<T>> small; // 小根堆，裡面存放大的值
    Heap<T, std::less<T>> big;      // 大根堆，裡面存放前k小的值
    //中位數就是big.top()
    DualHeap() {}
 
    void update() {
        if (big.size() == 0 and small.size() == 0) {
            return;
        }
 
        while (big.size() > small.size() + 1) {
            T x = big.top();
            big.pop();
            small.push(x);
        }
 
        while (big.size() < small.size()) {
            T x = small.top();
            small.pop();
            big.push(x);
        }
    }
 
    void push(T val) {
        if (big.size() == 0) {
            big.push(val);
            return;
        }
 
        if (val <= big.top()) {
            big.push(val);
        } else {
            small.push(val);
        }
 
        update();
    }
 
    void erase(T val) {
        assert(big.size() >= 1);
 
        if (val <= big.top()) {
            big.erase(val);
        } else {
            small.erase(val);
        }
 
        update();
    }

};
  • 可刪堆
template <class T, class Cmp = std::less<T>>
struct Heap {//可刪堆
    std::priority_queue<T, std::vector<T>, Cmp> qPush, qErase; // Heap=qPush-qErase
    i64 sum;
 
    Heap() : sum{0} {}
 
    void push(T x) {
        qPush.push(x);
    }
 
    void erase(T x) {
        qErase.push(x);
    }
 
    T top() {
        while (!qErase.empty() && qPush.top() == qErase.top())
            qPush.pop(), qErase.pop();
        return qPush.top();
    }
 
    void pop() {
        while (!qErase.empty() && qPush.top() == qErase.top()) {
            qPush.pop(), qErase.pop();
        }
 
        qPush.pop();
    }
 
    int size() {
        return qPush.size() - qErase.size();
    }
};

滑動視窗中位數

https://codeforces.com/gym/104901/problem/K

選區間內一個數，讓區間內每個數到這個數的距離之和最小，動態維護這個距離之和

首先一個經典結論，這個數就是中位數。

那麼問題轉化成：滑動視窗維護區間每個數到中位數的距離之和 \(ans\)。

顯然，ans 是所有比中位數 \(mid\) 與每個比中位數小的數 \(less\) 的差加上每個比中位數大的數 \(large\) 與中位數的差，也就是

\[ans = cnt_{less}\times mid - sum_{less} + sum_{large} - cnt_{large}\times mid \]

那麼我們就是要維護三個資訊：

\(less, large, mid\)

這很像對頂堆，less和large的相關資訊都完全可以對應上其中的大根堆和小根堆

單純的對頂堆，是不支援刪除的。

想要維護滑動視窗的中位數，就得結合可刪堆。

直接從應用入手，把兩個對頂堆改成可刪堆其中即可，而對 sum 的動態維護則在可刪堆的 push，pop，erase 操作中實現即可，非常直觀。

template <class T, class Cmp = std::less<T>>
struct Heap {//可刪堆
    std::priority_queue<T, std::vector<T>, Cmp> qPush, qErase; // Heap=qPush-qErase
    i64 sum;//維護出這個堆對應的sum
 
    Heap() : sum{0} {}
 
    void push(T x) {//加入的同時更新sum
        sum += x;
        qPush.push(x);
    }
 
    void erase(T x) {//刪除的同時更新sum
        sum -= x;
        qErase.push(x);
    }
 
    T top() {
        while (!qErase.empty() && qPush.top() == qErase.top())
            qPush.pop(), qErase.pop();
        return qPush.top();
    }
 
    void pop() {//一樣，更新sum
        while (!qErase.empty() && qPush.top() == qErase.top()) {
            qPush.pop(), qErase.pop();
        }
        sum -= qPush.top();
        qPush.pop();
    }
 
    int size() {//真實的個數也可以透過可刪堆維護出來
        return qPush.size() - qErase.size();
    }
};

template<class T>
struct DualHeap {//結合可刪堆，形成可刪對頂堆
 
    Heap<T, std::greater<T>> small; // small root
    Heap<T, std::less<T>> big;      // big root
 
    DualHeap() {}
 
    void update() {
        if (big.size() == 0 and small.size() == 0) {
            return;
        }
 
        while (big.size() > small.size() + 1) {
            T x = big.top();
            big.pop();
            small.push(x);
        }
 
        while (big.size() < small.size()) {
            T x = small.top();
            small.pop();
            big.push(x);
        }
    }
 
    void push(T val) {
        if (big.size() == 0) {
            big.push(val);
            return;
        }
 
        if (val <= big.top()) {
            big.push(val);
        } else {
            small.push(val);
        }
 
        update();
    }
 
    void erase(T val) {
        assert(big.size() >= 1);
 
        if (val <= big.top()) {
            big.erase(val);
        } else {
            small.erase(val);
        }
 
        update();
    }

    i64 getResult() {
        if (big.size() == 0) {return 0;}//說明只有一個數
 
        int x{big.top()}; 
 
        i64 ans1 = 1LL * x * big.size() - big.sum;//比中位數小的數的貢獻
        i64 ans2 = small.sum - 1LL * x * small.size();//比中位數大的數的貢獻
 
        return ans1 + ans2;
    }

};

void solve()
{
#define tests
    int n; i64 k; std::cin >> n >> k; std::vector<int> a(n); for (auto& ai : a) {std::cin >> ai; --ai;} 
    for (int i = 0; i < n; i++) {a[i] -= i;}
 
    DualHeap<int> dheap; int ans{}; for (int l = 0, r = 0; r < n; r++) {//這題固定右指標，移動左指標更方便
        dheap.push(a[r]);
        while (l <= r and dheap.getResult() > k) {//如果不滿足條件，就繼續移動左指標
            dheap.erase(a[l]);
            l += 1;
        }
        ans = std::max(ans, r - l + 1);
    }
 
    std::cout << ans << '\n';

}

相關文章