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;
}