const int MAXN = 2e5 + 10; int a[MAXN]; int val[MAXN]; #define mid ((l + r)>>1) int L[MAXN << 5]; int R[MAXN << 5]; int cnt[MAXN << 5]; ll sum[MAXN << 5]; int T[MAXN], tcnt; int iBuild(int l, int r) { int rt = ++tcnt; cnt[rt] = 0; sum[rt] = 0; if (l < r) { L[rt] = iBuild(l, mid); R[rt] = iBuild(mid + 1, r); } return rt; } // 在版本为pre的树的基础上,给x位置(val的离散排名)加上值val int iAdd(int pre, int l, int r, int x, int val) { int rt = ++tcnt; R[rt] = R[pre]; L[rt] = L[pre]; cnt[rt] = cnt[pre] + 1; sum[rt] = sum[pre] + val; if (l < r) { if (x <= mid) L[rt] = iAdd(L[pre], l, mid, x, val); else R[rt] = iAdd(R[pre], mid + 1, r, x, val); } return rt; } //查询(u,v]中排名为rk的数 int iVal(int u, int v, int l, int r, int rk) { //比整个区间的数都多,是INF if (rk > cnt[v] - cnt[u]) exit(-1); // 主席树上二分,找到rk所在的叶子 while (l < r) { int tmp = cnt[L[v]] - cnt[L[u]]; if (tmp < rk) { rk -= tmp; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } return val[l] ; } //查询(u,v]中值为va的数的排名 int iRnk(int u, int v, int l, int r, int va) { //比区间中最大的数都大 if (va > val[r]) return cnt[v] - cnt[u] + 1; // 主席树上二分,找到va所在的叶子 while (l < r) { if (val[mid] < va) u = R[u], v = R[v], l = mid + 1; else u = L[u], v = L[v], r = mid; } return l; } //查询(u,v]中排名不超过rk的数的个数 int iCntRnk(int u, int v, int l, int r, int rk) { int res = 0; // 主席树上二分,找到rk所在的叶子 while (l < r && rk < r) { if (mid <= rk) { //整个左子树都是<=rk,把整个左子树加上 res += cnt[L[v]] - cnt[L[u]]; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } // rk比最后找到的叶子还大,把最后的叶子加上 if (rk >= l) res += cnt[v] - cnt[u]; return res ; } //查询(u,v]中排名不超过rk的数的和 ll iSumRnk(int u, int v, int l, int r, int rk) { //上面的cnt变成sum ll res = 0; while (l < r && rk < r) { if (mid <= rk) { //整个左子树都是<=rk,把整个左子树加上 res += sum[L[v]] - sum[L[u]]; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } // rk比最后找到的叶子还大,把最后的叶子加上 if (rk >= l) res += sum[v] - sum[u]; return res ; } //查询(u,v]中值不超过va的数的个数 int iCntVal(int u, int v, int l, int r, int va) { int res = 0; while (l < r && va < val[r]) { if (val[mid] <= va) { //整个左子树都是<=va,把整个左子树加上 res += cnt[L[v]] - cnt[L[u]]; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } // va比最后找到的叶子还大,把最后的叶子加上 if (va >= val[l]) res += cnt[v] - cnt[u]; return res; } //查询(u,v]中值不超过va的数的和 ll iSumVal(int u, int v, int l, int r, int va) { //上面的cnt变成sum ll res = 0; while (l < r && va < val[r]) { if (val[mid] <= va) { //整个左子树都是<=va,把整个左子树加上 res += sum[L[v]] - sum[L[u]]; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } // va比最后找到的叶子还大,把最后的叶子加上 if (va >= val[l]) res += sum[v] - sum[u]; return res; } //查询(u,v]中,前k小的数的和 ll iSumK(int u, int v, int l, int r, int k) { //上面的cnt变成sum ll res = 0; while (l < r && k > 0) { if (cnt[L[u]] - cnt[L[v]] <= k) { //整个左子树不超过k个,把整个左子树加上 res += sum[L[v]] - sum[L[u]]; k -= cnt[L[u]] - cnt[L[v]]; u = R[u], v = R[v], l = mid + 1; } else u = L[u], v = L[v], r = mid; } // 走到叶子之后还有剩余的k,把这一部分加上(有足够的剩下吗?) res += 1LL * k * val[l]; return res; }
主席树还能求前k小的和,和“排名不超过rk的数的和”不同。
(也能求前k小的个数,但是这明显就是等于k)。