AtCoder Beginner Contest 234 D - Prefix K-th Max
問題の要約
(1,2,...,N)の順列P=(P1,P2,...,PN)、および正整数Kが与えられる。
i=K,K+1,...,Nについて、以下を求めよ。
・Pの先頭i項のうち、K番目に大きい値
制約
1<=K<=N<=5*10^5
入力
N K
P1 P2 ... PN
考え方
1,Piの値は(1,2,...,N)の順列なので重複はない。
2,K番目に大きい値ではなく、データ構造にK個の値だけを入れておけば、最小値に高速でアクセスできればよい。
こんかいはmapを使う。C++のmapは要素を取得する時間はO(logN)なので十分高速。
3,mapの要素数がK個になるように、削除と追加を繰り返して、最小値を出力していけばよい。
実際のプログラム
#include<iostream>
#include<vector>
#include<map>
using namespace std;
int main(){
long long N,K;
cin >> N >> K;
vector<long long> P(N);
for(long long i = 0; i < N; i++){
cin >> P[i];
}
map<long long, long long> mp;
for(long long i = 0; i < K; i++){
mp.insert(make_pair(P[i],i));
}
for(long long i = K; i < N; i++){
cout << mp.begin()->first << endl;
if(mp.begin()->first <= P[i]){
long long min_num = mp.begin()->first;
mp.erase(min_num);
mp.insert(make_pair(P[i],i));
}
}
cout << mp.begin()->first << endl;
return 0;
}