学习于https://atcoder.jp/contests/abc449/tasks/abc449_e
用于优化最后的查询
using namespace std;// ==================== 主席树模板（可持久化权值线段树） ====================
// 功能：区间第k小 / 第k大 / 区间小于x的数的个数 /
// 时间：建树 O(n log n)，查询 O(log n)
// 空间：O(n log n) 个节点struct ChairmanTree {// 节点结构struct Node {int l, r; // 左右子节点编号int sum; // 区间内的数字个数Node() : l(0), r(0), sum(0) {}};vector<Node> tree; // 节点池，tree[0]是空节点vector<int> root; // root[i]表示前i个元素对应的版本根节点vector<int> sorted; // 离散化数组int n; // 离散化后的值域大小int cnt; // 当前节点个数// 初始化，预估节点数 n * (log(n) + 2)ChairmanTree(int maxn = 200000) {tree.reserve(maxn * 20);tree.push_back(Node()); // tree[0]：空节点root.push_back(0); // root[0]：空版本cnt = 0;}// 离散化void init(vector<int>& a) {sorted = a;sort(sorted.begin(), sorted.end());sorted.erase(unique(sorted.begin(), sorted.end()), sorted.end());n = sorted.size();}// 获取x的离散化下标（1-based）inline int getId(int x) {return lower_bound(sorted.begin(), sorted.end(), x) - sorted.begin() + 1;}// 获取离散化前的原值inline int getValue(int id) {return sorted[id - 1];}// 新建节点，复制旧节点inline int clone(int old) {tree.push_back(tree[old]);return ++cnt;}// 在prev版本基础上，在位置pos插入一个数// 返回新版本根节点int update(int prev, int l, int r, int pos) {int cur = clone(prev);tree[cur].sum++;if (l != r) {int mid = (l + r) >> 1;if (pos <= mid) {int newLeft = update(tree[prev].l, l, mid, pos);tree[cur].l = newLeft;} else {int newRight = update(tree[prev].r, mid + 1, r, pos);tree[cur].r = newRight;}}return cur;}// 构建所有版本void build(vector<int>& a) {int m = a.size();root.resize(m + 1);root[0] = 0;for (int i = 0; i < m; i++) {int pos = getId(a[i]);root[i + 1] = update(root[i], 1, n, pos);}}// ==================== 查询接口 ====================// 1. 查询[L,R]区间内第k小的数（1-based，k>=1）// L,R：原数组下标（1-based），表示区间[L,R]int queryKth(int L, int R, int k) {// 前R个版本 - 前L-1个版本 = 区间[L,R]的信息return query(root[L - 1], root[R], 1, n, k);}// 2. 查询前k个元素中第q小的数（你的原始问题）int queryPrefix(int k, int q) {return query(root[0], root[k], 1, n, q);}// 3. 查询[L,R]区间内第k大的数int queryKthMax(int L, int R, int k) {int len = R - L + 1;return query(root[L - 1], root[R], 1, n, len - k + 1);}// 4. 查询[L,R]区间中小于x的数的个数int queryLess(int L, int R, int x) {int pos = lower_bound(sorted.begin(), sorted.end(), x) - sorted.begin();// pos是sorted中第一个>=x的位置，即小于x的有pos个if (pos == 0) return 0;return queryCount(root[L - 1], root[R], 1, n, 1, pos);}// 5. 查询[L,R]区间中大于x的数的个数int queryGreater(int L, int R, int x) {int pos = upper_bound(sorted.begin(), sorted.end(), x) - sorted.begin();// pos是sorted中第一个>x的位置，即<=x的有pos个if (pos >= n) return 0;return queryCount(root[L - 1], root[R], 1, n, pos + 1, n);}// 6. 查询[L,R]区间中等于x的数的个数int queryEqual(int L, int R, int x) {if (!binary_search(sorted.begin(), sorted.end(), x)) return 0;int pos = getId(x);return queryCount(root[L - 1], root[R], 1, n, pos, pos);}// ==================== 内部实现 ====================// 查询第k小（递归实现）int query(int u, int v, int l, int r, int k) {if (l == r) return getValue(l);int mid = (l + r) >> 1;int leftSum = tree[tree[v].l].sum - tree[tree[u].l].sum;if (k <= leftSum) {return query(tree[u].l, tree[v].l, l, mid, k);} else {return query(tree[u].r, tree[v].r, mid + 1, r, k - leftSum);}}// 查询区间[L,R]中值在[ql,qr]范围内的个数int queryCount(int u, int v, int l, int r, int ql, int qr) {if (ql > r || qr < l) return 0;if (ql <= l && r <= qr) return tree[v].sum - tree[u].sum;int mid = (l + r) >> 1;return queryCount(tree[u].l, tree[v].l, l, mid, ql, qr) +queryCount(tree[u].r, tree[v].r, mid + 1, r, ql, qr);}
};// ==================== 使用示例 ====================int main() {ios::sync_with_stdio(false);cin.tie(nullptr);// 示例数据vector<int> a = {7, 2, 9, 1, 5, 3, 8, 4, 6};int n = a.size();// 初始化主席树ChairmanTree ct(n);ct.init(a);ct.build(a);cout << "数组: ";for (int x : a) cout << x << " ";cout << "\n\n";// 1. 区间第k小cout << "区间[2,6]第3小: " << ct.queryKth(2, 6, 3) << "\n";// 区间[2,6] = {2,9,1,5,3}，排序{1,2,3,5,9}，第3小是3// 2. 前k个第q小（你的问题）cout << "前5个第3小: " << ct.queryPrefix(5, 3) << "\n";// 前5个 = {7,2,9,1,5}，排序{1,2,5,7,9}，第3小是5// 3. 区间第k大cout << "区间[3,8]第2大: " << ct.queryKthMax(3, 8, 2) << "\n";// 区间[3,8] = {9,1,5,3,8,4}，排序{1,3,4,5,8,9}，第2大是8// 4. 小于x的个数cout << "区间[1,9]小于5的个数: " << ct.queryLess(1, 9, 5) << "\n";// {7,2,9,1,5,3,8,4,6} 中小于5的有{2,1,3,4}，共4个// 5. 大于x的个数cout << "区间[1,9]大于5的个数: " << ct.queryGreater(1, 9, 5) << "\n";// {7,9,8,6}，共4个// 6. 等于x的个数cout << "区间[1,9]等于5的个数: " << ct.queryEqual(1, 9, 5) << "\n";// 只有1个return 0;
}```