第k大数查询 链表优化

Posted 2022-11-12 goto_1600

#include<bits/stdc++.h>
using namespace std;
const int N=500010;
int pre[N];
int ne[N];
int a[N];
int id[N];
int c[N];
int d[N];
int main()
	int n,k;
	scanf("%d%d",&n,&k);
	for(int i=1;i<=n;i++)	scanf("%d",&a[i]),id[a[i]]=i;
	long long res=0;
	set<int>s;
	
	for(int i=n;i>=1;i--)
	
		int l=0;
		int r=0;
		int ctl=0;
		int ctr=0;
		auto xx=s.lower_bound(id[i]);
		if(xx!=s.end())	r=*xx;
		if(s.size() and xx!=s.begin())	l=*(--xx);
		s.insert(id[i]);
		int j=r;
		d[0]=id[i];
		c[0]=id[i];
		for(;j && ctr<k;)
		
			d[++ctr]=j;
			j=ne[j];
		
		if(ctr<k)
			d[++ctr]=n+1;
		j=l;
		for(;j && ctl<k ;)
		
			c[++ctl]=j;
			j=pre[j];
		
		if(ctl<k)
			c[++ctl]=0;
		for(int j=k-1,t=0;j>=0;j--,t++)
		
			if(j+1<=ctl && t+1<=ctr )
			
				res+=1ll*i*(c[j]-c[j+1])*(d[t+1]-d[t]);
				
		
		
		if(r!=0)
			pre[r]=id[i];
		if(l!=0)
			ne[l]=id[i];
		pre[id[i]]=l;
		ne[id[i]]=r;
		// for(int g=0;g<=ctl;g++)
			// cout<<c[g]<<" ";
		// cout<<endl;
		// for(int g=0;g<=ctr;g++)
			// cout<<d[g]<<" ";
		// cout<<endl;
			// cout<<n-i+1<<" turn"<<res<<endl;
	
	cout<<res<<endl;
	
	return 0;

/*

5 2
1 2 3 4 5
*/