Luogu P2717 寒假作业(平衡树)

时间:2022-06-10 08:48:09

P2717 寒假作业

题意

题目背景

\(zzs\)和\(zzy\)正在被寒假作业折磨,然而他们有答案可以抄啊。

题目描述

他们共有\(n\)项寒假作业。\(zzy\)给每项寒假作业都定义了一个疲劳值\(A_i\),表示抄这个作业所要花的精力。\(zzs\)现在想要知道,有多少组连续的寒假作业的疲劳值的平均值不小于\(k\)?

简单地说,给定\(n\)个正整数\(A_1,A_2,A_3,\dots ,A_n\),求出有多少个连续的子序列的平均值不小于\(k\)。

输入输出格式

输入格式:

第一行两个正整数,\(n\)和\(k\)。

第二行到第\(n+1\)行,每行一个正整数\(A_i\)。

输出格式:

一个非负整数。

输入输出样例

输入样例#1:

3 2
1
2
3

输出样例#1:

4

说明

样例解释:共有\(6\)个连续的子序列,分别是\((1),(2),(3),(1,2),(2,3),(1,2,3)\),平均值分别为\(1,2,3,1.5,2.5,2\),其中平均值不小于\(k\)的共有\(4\)个。

对于\(20\%\)的数据,\(1\leq n\leq 100\);

对于\(50\%\)的数据,\(1\leq n\leq 5000\);

对于\(100\%\)的数据,\(1\leq n\leq 100000\);

对于\(100\%\)的数据,\(1\leq A_i\leq 10000,1\leq k\leq 10000\)。

思路

统计一个前缀和\(s\),然后考虑\([i,j]\)区间的平均值怎样才能不小于\(k\)呢?

\[\frac{s[j]-s[i-1]}{j-i+1}\geq k
\]

\[s[j]-k\times j\geq s[i-1]-(i-1)\times k
\]

诶,那对于每个\(j\)查询有多少个\(i\)满足上述式子不就好了吗?

在这里我选择了平衡树完成查询,平衡树选用的是\(fhq\ Treap\)。

还可以考虑的是,与处理出所有的\(s[i-1]-(i-1)\times k\),然后离散化弄到线段树上做,也是可行的,不过代码我就没有写了。

AC代码

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const LL MAXN=2e5+5;
LL n,k,cnt,rt,ans,a[MAXN],s[MAXN];
struct Node
{
LL val,rnd,sz;
LL ls,rs;
#define val(a) node[a].val
#define rnd(a) node[a].rnd
#define sz(a) node[a].sz
#define ls(a) node[a].ls
#define rs(a) node[a].rs
}node[MAXN];
LL read()
{
LL re=0;char ch=getchar();
while(!isdigit(ch)) ch=getchar();
while(isdigit(ch)) re=(re<<3)+(re<<1)+ch-'0',ch=getchar();
return re;
}
LL new_node(LL v)
{
val(++cnt)=v;
rnd(cnt)=rand();
sz(cnt)=1;
ls(cnt)=rs(cnt)=0;
return cnt;
}
void update(LL p){sz(p)=sz(ls(p))+sz(rs(p))+1;}
LL merge(LL x,LL y)
{
if(!x||!y) return x+y;
if(rnd(x)<rnd(y))
{
rs(x)=merge(rs(x),y);
update(x);
return x;
}
else
{
ls(y)=merge(x,ls(y));
update(y);
return y;
}
}
void split(LL now,LL v,LL &x,LL &y)
{
if(!now) x=y=0;
else
{
if(val(now)<=v)
{
x=now;
split(rs(now),v,rs(now),y);
}
else
{
y=now;
split(ls(now),v,x,ls(now));
}
update(now);
}
}
int main()
{
srand(time(0));
n=read(),k=read();
for(LL i=1;i<=n;i++) a[i]=read(),s[i]=s[i-1]+a[i];
for(LL i=1;i<=n;i++)
{
LL x,y;
split(rt,s[i]-k*i,x,y);
ans+=sz(x);
rt=merge(x,y);
split(rt,s[i-1]-k*(i-1),x,y);
rt=merge(merge(x,new_node(s[i-1]-k*(i-1))),y);
if(s[i]-k*i>=s[i-1]-k*(i-1)) ans++;
}
printf("%lld",ans);
return 0;
}