Official

E - Best Performances Editorial by en_translator


There are various solutions to this problem. This time, we will adopt an approach of using two instances of the same data structure.

Here, we use a multiset (a data structure that manages a multiset). Any data structure that supports the following operations can be used.

  • Storing a multiset. In other words, it can maintain multiple copies of the same value.
  • Accessing the maximum (or minimum) value (or both).
  • Removing a specific element.

Consider maintaining three states: a multiset \(X\) that stores the largest \(K\) values, another multiset \(Y\) that stores the smallest \(K\) values, and the sum \(s\) of \(X\).

\(X\) is initialized with the largest \(K\) values (this time, \(K\) copies of \(0\)s), \(Y\) with the remaining values in \(A\) (this time, \((N-K)\) copies of \(0\)), and \(s\) with the sum of \(X\).

We define the following operations: balance, add, and erase. This can be implemented with two multisets. During each operation, we also update \(s\) accordingly.

  • balance … if \(X\) does not have exactly \(K\) elements, repeatedly transfer the largest element of \(Y\) to \(X\) until it has exactly \(K\) elements. Then, while (\(X\)’s minimum value) \(<\) (\(Y\)’s maximum value), repeatedly swap them.
  • add … insert \(v\) to \(Y\), and do balance.
  • erase … remove a specific value \(v\). If \(v\) contains \(X\), remove it; otherwise, remove \(v\) from \(Y\). Then, do balance.

(Bonus: in this problem, the size of \(X\) never exceeds \(K\), but even if an update makes the size greater than \(K\), it can be still treated by transferring the smallest elements of \(X\) to \(Y\).)

When \(A_i\) is updated, it is sufficient to apply the following operations.

  • add new \(A_i\).
  • Then, erase original \(A_i\).

(You can do them other way round, but that will yield edge cases where the sum of the sizes of \(X\) and \(Y\) goes below \(0\) or both \(X\) and \(Y\) becomes empty when \(N=1\). In this problem, you can handle them by adding many \(0\)s as sentinels.)

In this problem, we can estimate that the number of swaps in balance is \(O(1)\), so the overall time complexity is \(O(Q \log N)\).

This solution can be adopted to other problems. When implementing, beware that \(X\) and \(Y\) may become an empty set.

Sample code (C++):

#include<bits/stdc++.h>

using namespace std;

int k;
multiset<int> x,y;
long long s;

void balance(){
  while(x.size()<k){
    auto iy=y.end();iy--;
    x.insert((*iy));
    s+=(*iy);
    y.erase(iy);
  }
  if(x.empty() || y.empty()){return;}

  while(1){
    auto ix=x.begin();
    auto iy=y.end();iy--;
    int ex=(*ix);
    int ey=(*iy);
    if(ex >= ey){break;}
    s+=(ey-ex);
    x.erase(ix);
    y.erase(iy);
    x.insert(ey);
    y.insert(ex);
  }
}

void add(int v){
  y.insert(v);
  balance();
}

void erase(int v){
  auto ix=x.find(v);
  if(ix!=x.end()){ s-=v; x.erase(ix); }
  else{ y.erase(y.find(v)); }
  balance();
}

int main(){
  int n;
  cin >> n >> k;
  vector<int> a(n,0);
  for(int i=0;i<k;i++){ x.insert(0); }
  for(int i=k;i<n;i++){ y.insert(0); }
  s=0;
  
  int q;
  cin >> q;
  while(q>0){
    q--;
    int p,w;
    cin >> p >> w;
    p--;
    add(w);
    erase(a[p]);
    a[p]=w;
    cout << s << "\n";
  }
  return 0;
}

posted:
last update: