【权值线段树】Codeforces Round #510 (Div. 2) D. Petya and Array

【权值线段树】Codeforces Round #510 (Div. 2) D. Petya and Array

https://codeforces.com/contest/1042/problem/D

1. 题意

给出一个长度为n的数组，问数组的有多少个区间和是小于t的？
$(1≤n≤200000,|t|≤2⋅10^{14})$(1≤n≤200000,∣t∣≤2⋅1014)

2. 思路

• 直接求肯定不行，时间复杂度太大。
• 可以先求前缀和，因为数据太大，所以前缀和可以离散化。
• 求完前缀和以后，可以从后往前将前缀和向线段树上更新，这样的话线段树上的节点可以记录下前缀和的个数（因为二叉树的左儿子的值小于父亲节点，右儿子大于父亲节点）。
• 所以计数的时候递归过程累计经过节点的左儿子的值，最后就能求出所以区间的解。注意，计数过程用的是sum[i]+t，而更新线段树过程用的是sum[i]
• 另外注意n=1的特殊情况

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const ll MAXN=2e5+100;

ll sum[MAXN];
vector<ll>vec;//离散化
//权值线段树
ll T[MAXN<<3];//线段树节点
#define lson rt<<1,l,mid
#define rson rt<<1|1,mid+1,r



int Find(ll x)
{
	return lower_bound(vec.begin(),vec.end(),x)-vec.begin()+1;
}
void build(int rt,int l,int r)
{
	T[rt]=0;
	if(l==r)
		return;
	int mid=(l+r)>>1;
	build(lson);
	build(rson);
}

void Update(int pos,ll val,int rt,int l,int r)
{
	T[rt]+=val;
	if(l==r)return;
	int mid=(l+r)>>1;
	if(pos<=mid)
		Update(pos,val,lson);
	else
		Update(pos,val,rson);
}


ll res=0;
void Query(int pos,int rt,int l,int r)
{

	if(l==r)
		return ;
	int mid=(l+r)>>1;
	if(pos<=mid)
		Query(pos,lson);
	else
	{
		res+=T[rt<<1];
		Query(pos,rson);
	}
	
}



ll a[MAXN];
int main()
{
	ll n,t;
	cin>>n>>t;

	for(ll i=1;i<=n;i++)
	{
		cin>>a[i];
		sum[i]=sum[i-1]+a[i];
		vec.push_back(sum[i]);
	}
	if(n==1)
    {
        if(a[1]<t)
            cout<<1<<endl;
        else cout<<0<<endl;
        return 0;

    }
	//离散化
	for(ll i=1;i<=n;i++)
	{
		vec.push_back(sum[i]+t);
	}
	sort(vec.begin(),vec.end());
	vec.erase(unique(vec.begin(),vec.end()),vec.end());

	//线段树过程
	ll len=vec.size();
	build(1,1,len);

	Update(Find(sum[n]),1,1,1,len);
	for(ll i=n-1;i>=0;i--)
	{
		ll x=sum[i]+t;
		Query(Find(x),1,1,len);
		if(i>0)
			Update(Find(sum[i]),1,1,1,len);
	}
	cout<<res<<endl;

	return 0;
}