AtCoder Regular Contest 028 B 特別賞

February 19, 2023

Pythonであればheap, C++であればsetを2つ使えばよい。新しい要素を追加する度に、一方に上位kの要素、他方にk+1以降の要素があるように、2つのset, heap間で要素を交換する。類似の問題をLeetCodeのHard modeで見かけたことがある。

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define as_int(a) static_cast<int>(a)

int main() {
  int n, k;
  cin >> n >> k;
  vector<int> xv(n);
  rep(i, n) cin >> xv[i];

  set<pair<int, int>> first, second;
  rep(i, n) {
    first.insert({xv[i], i + 1});
    if (i + 1 < k)
      continue;

    while (as_int(first.size()) > k) {
      second.insert(*first.rbegin());
      first.erase(prev(first.end()));
    }
    while (second.size() && first.rbegin()->first > second.begin()->first) {
      first.insert(*second.begin());
      second.insert(*first.rbegin());
      first.erase(prev(first.end()));
      second.erase(second.begin());
    }
    cout << first.rbegin()->second << endl;
  }

  return 0;
}