第 k 大查询(kth) 1s 128MB O2 加速

题目描述

给定一个 $1\sim n$ 的排列 $a_1,a_2,a_3,\dots ,a_n$，想知道所有满足 $r-l\ge k$ 的区间 $[l,r]$ 里 的数中第 $k$ 大的元素，输出它们的和。

输入格式

第一行两个整数 $n,k$。

接下来一行 $n$ 个整数 $a_1,a_2,a_3,\dots ,a_n$。

输出格式

输出一行，一个整数表示答案。

样例输入 1

5 2
1 2 3 4 5

样例输出 1

30

样例解释

其中区间 $[1,5],[2,5],[3,5],[4,5]$ 的第 $2$ 大为 $4$，区间 $[1,4],[2,4],[3,4]$ 的第 $2$ 大为 $3$，区间 $[1,3],[2,3]$ 的第 $2$ 大为 $2$，区间 $[1,2]$ 的第 $2$ 大为 $1$。

数据规模

共 $10$ 组数据。 测试点 $1,2$ 满足 $1\le n\le 100$。

测试点 $3,4$ 满足 $1\le n\le 10^3$。

测试点 $5,6$ 满足 $k=1$。

测试点 $7,8$ 满足 $1\le n\le 5\times 10^4$。

对于 $100\%$ 的数据，满足 $1\le n\le 5\times 10^5,1\le k\le 50$。

#include<bits/stdc++.h>
using namespace std;
int n,k,tl,tr,ll,lr,ps,vp,cnt;
long long ans=0;
int l[500006],r[500006],v[500006];
set<pair<int,int> >st;
int main(){
	scanf("%d %d",&n,&k);
	for(int i=0;i<=n+1;i++){
		if(i>0&&i<=n){
			scanf("%d",&v[i]);
			st.insert({v[i],i});
		}
		r[i]=i+1;
		l[i]=i-1;
	}
	for(auto i:st){
		vp=i.first;
		ps=i.second;
		tl=l[ps];
		tr=r[ps];
		ll=lr=ps;
		cnt=1;
		while(cnt<k&&ll!=r[0]){
			ll=l[ll];
			cnt++;
		}
		while(cnt<k&&lr!=l[n+1]){
			lr=r[lr];
			cnt++;
		}
		if(cnt<k)break;
		while(ll<=ps&&lr!=l[n+1]){
			ans+=1ll*vp*(ll-l[ll])*(r[lr]-lr);
			ll=r[ll];
			lr=r[lr];
		}
		r[tl]=tr;
		l[tr]=tl;
	}
	printf("%lld",ans);
	return 0;
}