主席树，区间第k小 / 第k大 / 区间小于x的数的个数 等

学习于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;
}```