【算法 - 数据结构】主席树

2021/4/10 12:30:40

本文主要是介绍【算法 - 数据结构】主席树,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

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)。



这篇关于【算法 - 数据结构】主席树的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!


扫一扫关注最新编程教程