[Blog] [Docs] [Code] [Slides] [About]

CodeForces 1328F - Make k Equal

前缀和

2020-04-10 19:27 CST

Problem

首先将所有的数计数去重;求出个数和总和的前缀和和后缀和。

每次计算将左边所有元素变为$val-1$和(如果需要)右边所有元素变为$val+1$后再取出$k - cnt[val]$个变为$val$所需要的步数,然后左右互换再算一次。

Solution

#include <iostream>
#include <algorithm>
#define _ \
  cin.sync_with_stdio(false); \
  cout.sync_with_stdio(false); \
  cin.tie(NULL), cout.tie(NULL);
using namespace std;
const long long LLINF = 0x3f3f3f3f3f3f3f3f;

class Item {
public:
  long long val, cnt, sum;
  Item() {
    val = cnt = sum = 0LL;
  };
};

int n = 0, k = 0, tot = 0;
long long a[200005] = {};
Item s[200005] = {};
Item pre[200005] = {};
Item suf[200005] = {};

void init() {
  cin >> n >> k;
  for (int i = 1; i <= n; ++i) {
    cin >> a[i];
  }
  sort(a + 1, a + n + 1);
  long long last = 0, cnt = 0;
  for (int i = 1; i <= n; ++i) {
    if (a[i] != last) {
      if (last) {
        ++tot;
        s[tot].val = last;
        s[tot].cnt = cnt;
        s[tot].sum = last * cnt;
      }
      last = a[i];
      cnt = 1;
    } else {
      ++cnt;
    }
  }
  ++tot;
  s[tot].val = last;
  s[tot].cnt = cnt;
  s[tot].sum = last * cnt;
  for (int i = 1; i <= tot; ++i) {
    pre[i].cnt = pre[i - 1].cnt + s[i].cnt;
    pre[i].sum = pre[i - 1].sum + s[i].sum;
  }
  for (int i = tot; i >= 1; --i) {
    suf[i].cnt = suf[i + 1].cnt + s[i].cnt;
    suf[i].sum = suf[i + 1].sum + s[i].sum;
  }
}

inline long long costL(int i) { return (s[i].val - 1) * (pre[i - 1].cnt) - pre[i - 1].sum; }
inline long long costR(int i) { return suf[i + 1].sum - (s[i].val + 1) * (suf[i + 1].cnt); }
long long solve() {
  long long ans = LLINF;
  for (int i = 1; i <= tot; ++i) {
    if (s[i].cnt >= k) {
      ans = 0LL;
      break;
    } else {
      long long cnt = k - s[i].cnt;
      {
        /* L to R */
        long long cntL = min(cnt, pre[i - 1].cnt);
        long long cntR = max(cnt - cntL, 0LL);
        long long tmp = costL(i);
        if (cntR > 0) {
          tmp += costR(i);
        }
        ans = min(ans, tmp + cnt);
      }
      {
        /* R to L */
        long long cntR = min(cnt, suf[i + 1].cnt);
        long long cntL = max(cnt - cntR, 0LL);
        long long tmp = costR(i);
        if (cntL > 0) {
          tmp += costL(i);
        }
        ans = min(ans, tmp + cnt);
      }
    }
  }
  return ans;
}

int main() { _
  init();
  printf("%lld\n", solve());
  return 0;
}
<EOF>