P1801 黑匣子

题意概括

维护一个初始为空的整数集合（黑匣子）。有两种操作：

1. ADD(x)：将一个整数 $x$ 加入集合。
2. GET：将一个计数器 $i$ 加一，然后查询并输出集合中第 $i$ 小的数。

输入以两个数组形式给出：$a$ 数组包含 $m$ 个要依次ADD的元素， $u$ 数组包含 $n$ 个GET操作的触发时机。具体来说，当第 $u_j$ 个元素（即 $a_{u_j}$）被加入集合后，立即执行一次GET操作。

数据范围：

• $1 \leq n, m \leq 2 \times 10^{5}$
• $|a_i| \leq 2 \times 10^{9}$
• $u$ 序列单调不降。

解题思路

这是一个动态查询第 $k$ 小值的问题。由于查询的 $k$ 是单调递增的，我们可以使用对顶堆（两个优先队列）来高效解决。

我们维护两个堆：

1. 一个大根堆 h1，用于存放集合中较小的一部分元素。
2. 一个小根堆 h2，用于存放集合中较大的一部分元素。

核心思想是，对于第 $k$ 次 GET 请求，我们调整两个堆，使得大根堆 h1 中恰好包含当前集合中最小的 $k$ 个元素。这样，h1 的堆顶元素 h1.top() 就是我们需要的答案，即第 $k$ 小的数。

我们始终保持 h1 中所有元素都小于等于 h2 中的所有元素。

• 添加元素 x：
  • 如果 h1 为空或 x 小于等于 h1 的堆顶，则将 x 加入 h1。
  • 否则，将 x 加入 h2。
• GET 操作：
  • 这是第 $k$ 次 GET 请求。
  • 我们需要让 h1 的大小调整为 $k$。
  • 如果 h1.size() > k，则不断将 h1 的堆顶元素移到 h2 中，直到 h1.size() == k。
  • 如果 h1.size() < k，则不断将 h2 的堆顶元素移到 h1 中，直到 h1.size() == k。
  • 完成调整后，h1.top() 即为所求的第 $k$ 小值。

此方法的时间复杂度为 $O((m+n)\log m)$，空间复杂度为 $O(m)$。

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int m, n;
    cin >> m >> n;

    vector<int> a(m);
    for (int i = 0; i < m; ++i) {
        cin >> a[i];
    }

    vector<int> u(n);
    for (int i = 0; i < n; ++i) {
        cin >> u[i];
    }

    // h1 是大根堆，存小数部分
    // h2 是小根堆，存大数部分
    priority_queue<int> h1;
    priority_queue<int, vector<int>, greater<int>> h2;

    int p = 0; // u 数组的指针
    int k = 0; // GET 命令的计数器，即题目中的 i

    for (int i = 0; i < m; ++i) {
        int x = a[i];

        // 维护 h1 的元素总是小于等于 h2 的元素
        if (h1.empty() || x <= h1.top()) {
            h1.push(x);
        } else {
            h2.push(x);
        }

        // 检查是否需要执行 GET 操作
        while (p < n && u[p] == i + 1) {
            k++; // 第 k 次 GET

            // 调整两个堆，使得 h1 的大小为 k
            while (h1.size() > k) {
                h2.push(h1.top());
                h1.pop();
            }
            while (h1.size() < k) {
                h1.push(h2.top());
                h2.pop();
            }
            
            // h1 的堆顶即为第 k 小的元素
            cout << h1.top() << "\n";
            p++;
        }
    }

    return 0;
}

解题思路 (权值线段树)

由于 ADD 操作中的数值范围很大（$|a_i| \leq 2 \times 10^{9}$），我们不能直接以数值为下标建立数据结构。但是，我们注意到元素的总数 $m$ 不超过 $2 \times 10^5$。这启发我们使用 离散化 的技巧。

1. 离散化: 我们将所有可能加入黑匣子的 $m$ 个数收集起来，进行排序和去重，得到一个最多包含 $m$ 个唯一值的升序数组 vls。任何一个原始数值都可以通过二分查找映射到它在 vls 中的下标（即它的 秩）。

2. 权值线段树: 我们建立一棵线段树，其维护的区间是离散化后的秩的范围 $[0, \text{sz}-1]$，其中 sz 是唯一值的数量。线段树的每个节点 t[i] 存储一个计数，表示当前黑匣子中有多少个数的 秩 落在节点 i所代表的区间内。

   • ADD(x) 操作:

     1. 找到 $x$ 离散化后的秩 r。
     2. 在线段树中对 r 位置进行单点更新，将其计数加一。这表示一个秩为 r 的数被加入了集合。时间复杂度为 $O(\log m)$。

   • GET 操作:

     1. 假设这是第 $k$ 次 GET 请求，我们需要找到集合中第 $k$ 小的数。
     2. 这等价于在线段树上查询第 $k$ 个被标记的位置。我们可以从根节点开始，利用每个节点的计数值来判断第 $k$ 小的数落在左子区间还是右子区间，并递归查找。
     3. 如果根节点的左孩子的计数值 cnt 大于等于 $k$，说明第 $k$ 小的数在左子树，我们递归进入左子树继续查找第 $k$ 小。
     4. 否则，说明第 $k$ 小的数在右子树，我们递归进入右子树查找第 $k - \text{cnt}$ 小。
     5. 当到达叶子节点时，该叶子节点对应的秩就是答案的秩。我们再通过 vls 数组将秩映射回原始值。查询时间复杂度为 $O(\log m)$。

整体时间复杂度为 $O((m+n)\log m)$，空间复杂度为 $O(m)$。

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

vector<int> t;   // 线段树
vector<int> vls; // 离散化后的值数组
int sz;

// 更新线段树，在 pos 位置增加计数
// c: 当前节点下标, l, r: 当前节点代表的区间, pos: 要更新的位置的秩
void upd(int c, int l, int r, int pos) {
    if (l == r) {
        t[c]++;
        return;
    }

    int mid = l + (r - l) / 2;

    if (pos <= mid) {
        upd(c * 2, l, mid, pos);
    } else {
        upd(c * 2 + 1, mid + 1, r, pos);
    }
    
    t[c] = t[c * 2] + t[c * 2 + 1];
}

// 查询第 k 小的数，返回其在 vls 中的秩(下标)
// c: 当前节点下标, l, r: 当前节点代表的区间, k: 要查询的排名
int qry(int c, int l, int r, int k) {
    if (l == r) {
        return l;
    }

    int mid = l + (r - l) / 2;
    int cnt = t[c * 2]; // 左子树中的元素数量

    if (k <= cnt) {
        return qry(c * 2, l, mid, k);
    } else {
        return qry(c * 2 + 1, mid + 1, r, k - cnt);
    }
}


int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    
    int m, n;
    cin >> m >> n;

    vector<int> a(m);
    for (int i = 0; i < m; ++i) {
        cin >> a[i];
    }

    vector<int> u(n);
    for (int i = 0; i < n; ++i) {
        cin >> u[i];
    }
    
    // 离散化
    vls = a;
    sort(vls.begin(), vls.end());
    vls.erase(unique(vls.begin(), vls.end()), vls.end());
    sz = vls.size();
    
    t.resize(sz * 4 + 4);
    
    int p = 0; // u 数组的指针
    int K = 0; // GET 命令的计数器

    for (int i = 0; i < m; ++i) {
        // 找到 a[i] 离散化后的秩
        int r = lower_bound(vls.begin(), vls.end(), a[i]) - vls.begin();
        upd(1, 0, sz - 1, r);

        // 处理在添加 a[i] 后触发的所有 GET 请求
        while (p < n && u[p] == i + 1) {
            K++;
            int res = qry(1, 0, sz - 1, K);
            cout << vls[res] << "\n";
            p++;
        }
    }

    return 0;
}