← 返回题目页

T0568 算法七_期中_06(平衡树 弃用)

标程

#include <bits/stdc++.h>
using namespace std;

/*
黑匣子:
  1 x  插入 x
  2 x  删除一个 x(不存在则忽略)
  3 k  第 k 小(不足 k 个输出 -1)
  4 t  求 sum |x - t|

离线把所有出现过的数离散化,用树状数组维护「个数」和「数值和」两份前缀信息:
  - 第 k 小:树状数组上倍增找前缀和刚好 >= k 的位置。
  - sum|x-t|:设 <= t 的元素有 cnt 个、和为 sum,
      答案 = (sum - t*cnt) + (t*(tot-cnt) - (totSum - sum))
            = totSum - 2*sum + t*(2*cnt - tot)。
*/

struct BIT {
    int n;
    vector<int> cnt;
    vector<long long> sum;
    BIT(int n) : n(n), cnt(n + 1, 0), sum(n + 1, 0) {}
    void add(int i, int dc, long long ds) {
        for (; i <= n; i += i & -i) {
            cnt[i] += dc;
            sum[i] += ds;
        }
    }
    // 前缀 [1..i] 的个数与数值和
    void query(int i, int &c, long long &s) {
        c = 0; s = 0;
        for (; i > 0; i -= i & -i) {
            c += cnt[i];
            s += sum[i];
        }
    }
    // 第 k 小(1-based),返回离散化后的下标;不存在返回 -1
    int kth(int k, int total) const {
        if (k > total) return -1;
        int pos = 0;
        int step = 1;
        while ((step << 1) <= n) step <<= 1;
        for (; step; step >>= 1) {
            int nxt = pos + step;
            if (nxt <= n && cnt[nxt] < k) {
                pos = nxt;
                k -= cnt[nxt];
            }
        }
        return pos + 1;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int m;
    cin >> m;
    vector<int> op(m), arg(m);
    vector<long long> vals;
    for (int i = 0; i < m; i++) {
        cin >> op[i] >> arg[i];
        if (op[i] == 1 || op[i] == 2) vals.push_back(arg[i]);
    }
    sort(vals.begin(), vals.end());
    vals.erase(unique(vals.begin(), vals.end()), vals.end());

    auto idOf = [&](int x) {
        return int(lower_bound(vals.begin(), vals.end(), (long long)x) - vals.begin()) + 1;
    };

    BIT bit((int)vals.size());
    vector<int> freq(vals.size() + 1, 0);
    int total = 0;
    long long totSum = 0;

    for (int i = 0; i < m; i++) {
        int o = op[i], x = arg[i];
        if (o == 1) {
            int id = idOf(x);
            bit.add(id, 1, (long long)x);
            freq[id]++;
            total++;
            totSum += x;
        } else if (o == 2) {
            int id = idOf(x);
            if (freq[id] > 0) {
                bit.add(id, -1, -(long long)x);
                freq[id]--;
                total--;
                totSum -= x;
            }
        } else if (o == 3) {
            int id = bit.kth(x, total);
            if (id == -1) cout << -1 << "\n";
            else cout << vals[id - 1] << "\n";
        } else {
            // sum |x - t|
            int cntLe = 0;
            long long sumLe = 0;
            int idx = int(lower_bound(vals.begin(), vals.end(), (long long)x) - vals.begin());
            // vals[0..idx-1] 都严格小于 x
            if (idx > 0) bit.query(idx, cntLe, sumLe);
            long long ans = (totSum - 2 * sumLe) + (long long)x * (2 * cntLe - total);
            cout << ans << "\n";
        }
    }
    return 0;
}

动画演示