少年，想學帶修改主席樹嗎 | BZOJ1901 帶修改區間第k小

有一道題（BZOJ 1901）是這樣的：n個數，m個詢問，詢問有兩種：修改某個數/詢問區間第k小。

不帶修改的區間第k小用主席樹很好寫，不會的同學可以看一下這個。

加上修改怎麽做呢？我們可以用數學老師成天講的類比思想：

可以發現，不修改的區間k小問題中，每加入一個原序列中的數，對應的主席樹在上一個的基礎上進行修改，而查詢的時候用右端點主席樹減去左端點左邊的主席樹。這樣的操作就像是維護前綴和：每次加入一個元素的時候，sum[i] = sum[i - 1] + a[i]；詢問的時候，則是sum[r] - sum[l - 1]。

sum數組可以用來不帶修改求前綴和，那麽假如我們要求帶修改的前綴和呢？樹狀數組可以做到。每次在位置p加入一個元素x的時候，對樹狀數組中每個 {p, p + (p & -p), ...} 都加上x；詢問前綴和的時候，則是求樹狀數組中每個 {p, p - (p & -p), ...} 的和。

現在我們來看看能否將樹狀數組和主席樹結合起來，實現動態修改。

一開始，樹狀數組的每個位置都對應著一棵主席樹——雖然實際上它們對應的都是同一棵空空蕩蕩的主席樹。

然後我們該往裏面加數了：當在位置p加入一個數x的時候，對樹狀數組中位置 {p, p + (p & -p), ...} 上的主席樹都進行“加入數x”的操作，也就是把數x對應的位置++。

然後可以處理詢問了。對於查詢，可以采用非遞歸的方式：維護兩個cur數組，分別記錄左端點（的左邊）對應的主席樹 {l, l - (l & -l), ...} 上的當前節點和右端點對應主席樹 {r, r - (r & -r), ...} 上的當前節點。每次對兩個cur數組上記錄的節點的data求和，然後相減，可以得到區間內的data。這個操作實際上就是樹狀數組求前綴和的操作。

詢問的代碼如下（寫得比較臃腫……但是意思應該是比較清楚的了）：

int query(int ql, int qr, int k){
    int l = 1, r = idx;
    for(int p = ql; p; p -= p & -p) cur1[p] = root[p];
    for(int p = qr; p; p -= p & -p) cur2[p] = root[p];
    while(l < r){
        int mid = (l + r) >> 1, sum1 = 0, sum2 = 0;
        for(int
 
 p = ql; p; p -= p & -p) sum1 += data[ls[cur1[p]]];
        for(int p = qr; p; p -= p & -p) sum2 += data[ls[cur2[p]]];
        if(sum2 - sum1 >= k){
            for(int p = ql; p; p -= p & -p) cur1[p] = ls[cur1[p]];
            for(int p = qr; p; p -= p & -p) cur2[p] = ls[cur2[p]];
            r = mid;
        }
        else{
            l = mid + 1, k -= sum2 - sum1;
            for(int p = ql; p; p -= p & -p) cur1[p] = rs[cur1[p]];
            for(int p = qr; p; p -= p & -p) cur2[p] = rs[cur2[p]];
        }
    }
    return lst[l];
}

那麽對於修改操作呢？當然是和初始化時加入每個數時一樣，對於樹狀數組中每個 {p, p + (p & -p), ...} 位置上的主席樹都進行修改咯。

void change(int old, int &k, int l, int r, int p, int x){
    k = ++tot;
    data[k] = data[old] + x, ls[k] = ls[old], rs[k] = rs[old];
    if(l == r) return;
    int mid = (l + r) >> 1;
    if(p <= mid) change(ls[old], ls[k], l, mid, p, x);
    else change(rs[old], rs[k], mid + 1, r, p, x);
}
void add(int p, int num, int x){
    while(p <= n) change(root[p], root[p], 1, idx, num, x), p += p & -p;
}

那麽這道題就做完啦。

下面是完整的代碼，加上近三十行的讀入優化後也只有100行，比起學長讓我近期寫的各種樹套樹……應該算是比較短小好寫的吧。

樹狀數組天下第一！

#include <cstdio>
#include <cmath>
#include <cstring>
#include <algorithm>
#define space putchar(' ')
#define enter putchar('\n')
using namespace std;
typedef long long ll;
template <class T>
void read(T &x){
    char c;
    bool op = 0;
    while(c = getchar(), c < '0' || c > '9')
        if(c == '-') op = 1;
    x = c - '0';
    while(c = getchar(), c >= '0' && c <= '9')
        x = x * 10 + c - '0';
    if(op) x = -x;
}
template <class T>
void write(T x){
    if(x < 0) putchar('-'), x = -x;
    if(x >= 10) write(x / 10);
    putchar('0' + x % 10);
}
const int N = 100005, M = 6000005;
int n, m, a[N], lst[N], idx;
int tot, root[N], data[M], ls[M], rs[M], cur1[N], cur2[N];
int qtype[N], q1[N], q2[N], q3[N];
void build(int &k, int l, int r){
    k = ++tot;
    if(l == r) return;
    int mid = (l + r) >> 1;
    build(ls[k], l, mid);
    build(rs[k], mid + 1, r);
}
void change(int old, int &k, int l, int r, int p, int x){
    k = ++tot;
    data[k] = data[old] + x, ls[k] = ls[old], rs[k] = rs[old];
    if(l == r) return;
    int mid = (l + r) >> 1;
    if(p <= mid) change(ls[old], ls[k], l, mid, p, x);
    else change(rs[old], rs[k], mid + 1, r, p, x);
}
void add(int p, int num, int x){
    while(p <= n) change(root[p], root[p], 1, idx, num, x), p += p & -p;
}
int query(int ql, int qr, int k){
    int l = 1, r = idx;
    for(int p = ql; p; p -= p & -p) cur1[p] = root[p];
    for(int p = qr; p; p -= p & -p) cur2[p] = root[p];
    while(l < r){
        int mid = (l + r) >> 1, sum1 = 0, sum2 = 0;
        for(int p = ql; p; p -= p & -p) sum1 += data[ls[cur1[p]]];
        for(int p = qr; p; p -= p & -p) sum2 += data[ls[cur2[p]]];
        if(sum2 - sum1 >= k){
            for(int p = ql; p; p -= p & -p) cur1[p] = ls[cur1[p]];
            for(int p = qr; p; p -= p & -p) cur2[p] = ls[cur2[p]];
            r = mid;
        }
        else{
            l = mid + 1, k -= sum2 - sum1;
            for(int p = ql; p; p -= p & -p) cur1[p] = rs[cur1[p]];
            for(int p = qr; p; p -= p & -p) cur2[p] = rs[cur2[p]];
        }
    }
    return lst[l];
}
int getpos(int x){
    return lower_bound(lst + 1, lst + idx + 1, x) - lst;
}
bool isQ(){
    char c;
    while(c = getchar(), c != 'Q' && c != 'C');
    return c == 'Q';
}
int main(){
    read(n), read(m), idx = n;
    for(int i = 1; i <= n; i++)
        read(a[i]), lst[i] = a[i];
    for(int i = 1; i <= m; i++){
        qtype[i] = isQ(), read(q1[i]), read(q2[i]);
        if(qtype[i]) read(q3[i]);
        else lst[++idx] = q2[i];
    }
    sort(lst + 1, lst + idx + 1);
    idx = unique(lst + 1, lst + idx + 1) - lst - 1;
    build(root[0], 1, idx);
    for(int i = 1; i <= n; i++) root[i] = root[0];
    for(int i = 1; i <= n; i++) add(i, getpos(a[i]), 1);
    for(int i = 1; i <= m; i++){
        if(qtype[i]) write(query(q1[i] - 1, q2[i], q3[i])), enter;
        else{
            add(q1[i], getpos(a[q1[i]]), -1);
            a[q1[i]] = q2[i];
            add(q1[i], getpos(a[q1[i]]), 1);
        }
    }
    return 0;
}

少年，想學帶修改主席樹嗎 | BZOJ1901 帶修改區間第k小