GYM102770E Easy DP Problem（可持久化线段树）

题意：

区间前K大树的和，用可持久化线段树完成，比赛的时候WA了好几发，这方面还是不够熟练。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+100;
const int M=maxn*40;
int n,m,q;
int a[maxn];//原始数组
int t[maxn];//离散化数组
int T[maxn];//第i棵线段树的根节点编号
int lson[M];
int rson[M];
int c[M];
ll sum[M];
int tot;
 
int build (int l,int r) {
    int root=tot++;
    c[root]=0;
    sum[root]=0;
    if (l!=r) {
        int mid=(l+r)>>1;
        lson[root]=build(l,mid);
        rson[root]=build(mid+1,r);
    }
    return root;
} 
int up (int root,int p,int v) {
    int newRoot=tot++;
    int tmp=newRoot;
    int l=1,r=m;
    while (l<r) {
        int mid=(l+r)>>1;
        if (p<=mid) {
            lson[newRoot]=tot++;
            rson[newRoot]=rson[root];
            newRoot=lson[newRoot];
            root=lson[root];
            r=mid;
        }
        else {
            rson[newRoot]=tot++;
            lson[newRoot]=lson[root];
            newRoot=rson[newRoot];
            root=rson[root];
            l=mid+1;
        }
        c[newRoot]=c[root]+v;
        sum[newRoot]=sum[root]+t[p]*v;
    }
    return tmp;
}
int query (int left_root,int right_root,int k) {
    int l=1,r=m;
    while (l<r) {
        int mid=(l+r)>>1;
        if (c[rson[left_root]]-c[rson[right_root]]>=k) {
            l=mid+1;
            left_root=rson[left_root];
            right_root=rson[right_root];
        }
        else {
            r=mid;
            k-=c[rson[left_root]]-c[rson[right_root]];
            left_root=lson[left_root];
            right_root=lson[right_root];
        }
    }
    return l;
}
 
int q1 (int left_root,int right_root,int l,int r,int L,int R) {
    //查询数量
    //l r为当前区间
    //L R为查询区间
    if (l>=L&&r<=R) return c[left_root]-c[right_root];
    int mid=(l+r)>>1;
    int ans=0;
    if (L<=mid) ans+=q1(lson[left_root],lson[right_root],l,mid,L,R);
    if (R>mid) ans+=q1(rson[left_root],rson[right_root],mid+1,r,L,R);
    return ans;  
}
 
ll q2 (int left_root,int right_root,int l,int r,int L,int R) {
    //查询数量
    //l r为当前区间
    //L R为查询区间
    if (l>=L&&r<=R) return sum[left_root]-sum[right_root];
    int mid=(l+r)>>1;
    ll ans=0;
    if (L<=mid) ans+=q2(lson[left_root],lson[right_root],l,mid,L,R);
    if (R>mid) ans+=q2(rson[left_root],rson[right_root],mid+1,r,L,R);
    return ans;  
}
ll in[maxn];
int main () {
    int _;
    for (ll i=1;i<maxn;i++) in[i]=in[i-1]+i*i;
    scanf("%d",&_);
    while (_--) {
        scanf("%d",&n);
        for (int i=1;i<=n;i++) scanf("%d",a+i),t[i]=a[i];
        sort(t+1,t+n+1);
        m=unique(t+1,t+n+1)-t-1;
        for (int i=1;i<=n;i++) a[i]=upper_bound(t+1,t+m+1,a[i])-t-1;
        tot=0;
        T[n+1]=build(1,m);
        for (int i=n;i;i--) T[i]=up(T[i+1],a[i],1);
        scanf("%d",&q);
        while (q--) {
            int l,r,k;
            scanf("%d%d%d",&l,&r,&k);
            int L=query(T[l],T[r+1],k);
            int R=m;
            //printf("%d\n",L);
            printf("%lld\n",q2(T[l],T[r+1],1,m,L+1,R)+(long long)(k-q1(T[l],T[r+1],1,m,L+1,R))*t[L]+in[r-l+1]);
        }
    }
}