PKU 2104 K-th Number

問題文→2104 -- K-th Number

sort()とか使うと間に合わないので　区間ごとのx以下の個数を二分探索で数えます。しかし二分探索にはソートが必要なのでマージソートの要領でセグメント木を作っておき、必要な区間に対してのみ二分探索を行えば間に合うみたいです。

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<vector>
#include<string>
#include<cmath>
#include<map>
#include<cstdlib>
#include<climits>
#include<cstring>
#include<queue>
#include<deque>
#define FOR(i,a,b) for(int i=(a);i<(b);++i)
#define REP(i,n)  FOR(i,0,n)
#define INF 99999999
#define EPS 0.0001
using namespace std;
typedef vector<string> VS;
const int ST_SIZE=(1<<18)-1;
const int MAX_N=100000;
const int MAX_M=5000;
int N,M;
int A[MAX_N];
int I[MAX_M],J[MAX_M],K[MAX_M];
int nums[MAX_N];
vector<int>dat[ST_SIZE];

void init(int k,int l,int r){
  if(r-l==1) dat[k].push_back(A[l]);
  else{
    int lch=k*2+1,rch=k*2+2;
    init(lch,l,(l+r)/2);
    init(rch,(l+r)/2,r);
    dat[k].resize(r-l);
    merge(dat[lch].begin(),dat[lch].end(),dat[rch].begin(),dat[rch].end(),dat[k].begin());
  }
}
int query(int i,int j, int x,int k,int l,int r){
  if(j<=l || r<=i) return 0;
  else if(i<=l && r<=j) return upper_bound(dat[k].begin(),dat[k].end(),x)-dat[k].begin();
  else{
    int lc=query(i,j,x,k*2+1,l,(l+r)/2);
    int rc=query(i,j,x,k*2+2,(l+r)/2,r);
    return lc+rc;
  }  
}
void solve(){
  for(int i=0;i<N;i++) nums[i]=A[i];  
  sort(nums,nums+N);
  
  init(0,0,N);
  
  for(int i=0;i<M;i++){
    int l=I[i],r=J[i]+1,k=K[i];
    int lb=-1,ub=N-1;
    while(ub-lb>1){
      int md=(ub+lb)/2;
      int c=query(l,r,nums[md],0,0,N);
      if(c>=k) ub=md;
      else lb=md;
    }
    cout<<nums[ub]<<endl;
  }
}
int main(){
  cin>>N>>M;
  REP(i,N) cin>>A[i];
  REP(i,M){ cin>>I[i]>>J[i]>>K[i]; I[i]--,J[i]--;}
  solve();
}