HDU 4217 Data Structure?(线段树)

#include <iostream>
#include <cstdio>
using namespace std;
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
const int maxn=262154;
int segs[maxn<<2],n,k;
void pushUp(int rt){
segs[rt]=segs[rt<<1]+segs[rt<<1|1];
}
void build(int l,int r,int rt){
if(l==r){
segs[rt]=1;
return;
}
int mid=(l+r)>>1;
build(lson),build(rson);
pushUp(rt);
}
int query(int p,int l,int r,int rt){
if(l==r){
segs[rt]=0;
return l;
}
int mid=(l+r)>>1,ret=0;
if(p<=segs[rt<<1])ret=query(p,lson);
else ret=query(p-segs[rt<<1],rson);
pushUp(rt);
return ret;
}

int main(){
int t,cas=1;
scanf("%d",&t);
while (t--)
{
scanf("%d%d",&n,&k);

long long ans=0;

build(1,n,1);
while (k--)
{
int ki,t=0;
scanf("%d",&ki);
ans+=query(ki,1,n,1);
}
printf("Case %d: %I64d\n",cas++,ans);
}
return 0;
}