Data Structure?

题目链接:https://vjudge.net/problem/HDU-4217

思路

主要就是这段话:
以查询第k大为例,权值线段树的核心是到每个结点,如果右子树的权值总和大于了k,则说明其第k大值在右子树,递归进入右子树。反之则说明第k大值在左子树。

特别注意:若要进入左子树,需要k减去右子树的总和,比如要找的元素是第5大,右子树权值总和为3,则需5-3=2,说明该节点的第5大值存在于左子树的第2大值中。那么从左子树递归下去,直到递归到一个数,那就是答案了。

代码

#include<bits/stdc++.h>
#define ios ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0)
#define debug  freopen("in.txt","r",stdin),freopen("out.txt","w",stdout);
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int maxn = 262144+10;
using namespace std;

int T,N,K;
struct node
{
    int l,r,cnt;
}tr[maxn*4];

void pushup(int u){
    tr[u].cnt = tr[u*2].cnt + tr[u*2+1].cnt;
}
void build(int l,int r,int u = 1){
    tr[u] = {l,r,1};
    if(l == r) return;
    int mid = (l+r)>>1;
    build(l,mid,u*2);
    build(mid+1,r,u*2+1);
    pushup(u);
}
void modify(int idx,int v,int u = 1){
    if(tr[u].l == idx && tr[u].r == idx) tr[u].cnt = v;
    else{
        int mid = (tr[u].l + tr[u].r)>>1;
        if(idx<=mid) modify(idx,v,u*2);
        else modify(idx,v,u*2+1);
        pushup(u);
    }
}
int query(int k,int u = 1){
    if(tr[u].l == tr[u].r && tr[u].cnt == k) return tr[u].l;
    if(tr[u*2].cnt >=k) return query(k,u*2);
    else return query(k-tr[u*2].cnt,u*2+1);
}

int main(){
    int kase = 0;
    cin>>T;
    while(T--){
        scanf("%d %d",&N,&K);
        build(1,N);
        ll ans = 0;
        while(K--){
            int ki;scanf("%d",&ki);
            int v = query(ki);
            ans += v;
            modify(v,0);
        }
        printf("Case %d: %lld\n",++kase,ans);
    }


    return 0;
}