The 17th Zhejiang Provincial Contest E. Easy DP Problem（主席树）

题目描述

题目链接

题目大意

给你一个序列a[]，以及q个查询。
对于每个查询，给出l，r，k，要求以a[l-r]作为b[]数组求出 dp[r-l+1][k]的值。

题目分析

通过分析题目给出的状态转移方程，我们可以发现：
$dp[r-l+1][k]=\sum _{i=1}^{r-l+1}{i}^2+b数组中前k大的数之和。$

$结果的前一部分我们可以用平方和公式求出：\sum _{i=1}^n{i}^2=n\ast (n+1)\ast (2\ast n+1)/6.$

后一部分的操作是：求任意区间[l,r]中的前k大的数之和。我们可以通过主席树来快速求出结果。

代码如下

#include <iostream>
#include <cmath>
#include <cstdio>
#include <set>
#include <string>
#include <cstring>
#include <map>
#include <algorithm>
#include <stack>
#include <queue>
#include <bitset>
#define LL long long
#define ULL unsigned long long
#define PII pair<LL,LL>
#define PDD pair<double,double>
#define x first
#define y second
using namespace std;
const int N=1e5+5,INF=1e9+7;
struct Node{
	int l,r;
	int cnt;			//维护当前段上的节点个数
	LL sum;				//当前段上的数值之和
}tr[N*21];
int root[N],idx;
LL a[N];
vector<LL> num;
int insert(int p,int l,int r,int x)			//基于p版本，向主席树中插入一个数
{
	int u=++idx;
	tr[u]=tr[p];
	tr[u].cnt++,tr[u].sum=tr[p].sum+num[x];		//sum维护的是离散化之前的数值和
	if(l==r) return u;
	int mid=l+r>>1;
	if(mid>=x) tr[u].l=insert(tr[p].l,l,mid,x);
	else tr[u].r=insert(tr[p].r,mid+1,r,x);
	return u;
}
LL query(int p,int q,int l,int r,int k)			//查找[l,r]区间内前k大的数之和
{
	if(l==r) return num[l]*k;
	int mid=l+r>>1;
	int size=tr[tr[p].r].cnt-tr[tr[q].r].cnt;	//因为是前k大，因此要先往右找
	if(size>=k) return query(tr[p].r,tr[q].r,mid+1,r,k);
	return query(tr[p].l,tr[q].l,l,mid,k-size)+tr[tr[p].r].sum-tr[tr[q].r].sum; 
}
int main()
{
	int t;
	scanf("%d",&t);
	while(t--)
	{
		idx=0;
		num.clear();
		memset(tr,0,sizeof tr); 
		memset(root,0,sizeof root);
		int n,q;
		scanf("%d",&n);
		for(int i=1;i<=n;i++)
		{
			scanf("%lld",&a[i]);
			num.push_back(a[i]);
		}
		sort(num.begin(),num.end());				//a[]的数据范围比较大，因此我们需要离散化
		num.erase(unique(num.begin(),num.end()),num.end());
		for(int i=1;i<=n;i++)
			a[i]=lower_bound(num.begin(),num.end(),a[i])-num.begin();
		
		for(int i=1;i<=n;i++)						//建树
			root[i]=insert(root[i-1],0,num.size()-1,a[i]);
		scanf("%d",&q);
		while(q--)
		{
			int l,r,k;
			scanf("%d%d%d",&l,&r,&k);
			LL t=r-l+1;
			printf("%lld\n",query(root[r],root[l-1],0,num.size()-1,k)+t*(t+1)*(2*t+1)/6);
		}
	}
	return 0;
}