对顶堆

对顶堆是一种利用一个大根堆和一个小根堆组合而成的数据结构。两个堆顶形成了一种沙漏结构,故称为对顶堆,常用于\(O(1)\)时间复杂度获取数据流的中位数。

class TwoHeaps {
public:
    priority_queue<int, vector<int>, less<int>> pq1;
    priority_queue<int, vector<int>, greater<int>> pq2;

    TwoHeaps() {}

    void add(int x) {
        if (pq1.empty() && pq2.empty()) {
            pq1.push(x);
        } else if (!pq1.empty() && x < pq1.top()) {
            pq1.push(x);
        } else {
            pq2.push(x);
        }
        adjustEqualSize();
    }

    double getMedian() {
        if (pq1.size() == pq2.size()) {
            return (pq1.top() + pq2.top()) / 2.0;
        } else {
            return (double) pq2.top();
        }
    }

private:
    void adjustEqualSize() {
        while (pq1.size() > pq2.size()) {
            pq2.push(pq1.top());
            pq1.pop();
        }
        while (pq2.size() > pq1.size() + 1) {
            pq1.push(pq2.top());
            pq2.pop();
        }
    }
};