競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 281E問題

ABC281E

解法

\(M\) 個の元を持っておき,小さい方から \(K\) 個の元の和も同時に保持しておく. 追加や削除をしたときに,一部の元しか更新されないので,それを利用して高速化する. つまり,差分を高速に更新出来ればよい.

実装

小さい方から先頭 \(K\) 個 \(front\) と, 残り \(M-K\) 個 \(back\) を分けて,それぞれ multiset で保持する. 同時に,\(front\) の和も保持する. 新しく元を \(x\) 追加するとき,\(front\) の一番後ろの元と, \(back\) の先頭の元を見れば,どちらに追加すれば良いのか分かる. 削除するときは,重複するときはどちらから除いても良い.
追加や削除をしたとき,\(front\) が \(K\) 個を超えたり, 足りなくなったりする. そのときに,\(back\) の方から補充したり渡したりして \(front\) が \(K\) 個になるようにする関数を用意する. また,\(front\) に追加や削除をするときは, 同時に和も更新する.

使っている記号,マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

struct D{
  void front_insert(ll p){
    front.insert(p);
    tot += p;
  }
 
  void front_erase(auto it){
    front.erase(it);
    tot -= *it;
  }
 
public:
  ll tot; // total of front
  ll k; // m は使わない
  multiset<ll> front, back; // index は持たない代わりに,multiset

  D(){}
  D(ll k): k(k){
    tot = 0;
  }

  // rem: front に insert, erase するときは,tot も一緒に更新する.
  // back のときは tot の更新が不要.
  // 更新し忘れが心配なら,push_front, erase_front の様な関数を作る.

  void adjust(){
    while(front.size() < k && back.size()){
      auto it = back.begin();
      front_insert(*it);
      back.erase(it);
    }
   
    while(front.size() > k){
      auto it = prev(front.end());
      front_erase(it);
      back.insert(*it);
    }
 
  }

  void push(ll x){
    front_insert(x);
    adjust();
  }

  void pop(ll x){
    auto it = front.lower_bound(x);
    if(it != front.end()) { // front
      front_erase(it);
    }else{
      it = back.lower_bound(x);
      if(it != back.end()){ // back
        back.erase(it);
      }else{
        assert(false);
      }
    }

    adjust();
  }

};

int main() {
  ll n,m,k;
  cin >> n >> m >> k;
  vll a(n); rep(i,n) { cin >> a[i]; }
  D d(k);

  rep(i,m-1) {
    d.push(a[i]);
  }
  rep(i,n-m+1){
    d.push(a[i+m-1]);
    cout << d.tot << endl;
    d.pop(a[i]);
  }



  return 0;
}