CodeForces1328F 贪心

题意

给你一个由n个元素和k≤n的整数组成的数组a。 您要在数组a中获得至少k个相等的元素。一次移动，可以执行以下两个操作之一： 取数组的最小元素之一，并将其值增加1 取数组的最大元素之一，并将其值减少1 您的任务是计算在数组中至少获得k个相等元素所需的最小移动次数。

Input

输入的第一行包含两个整数n和k（1≤k≤n≤2⋅105），即a中的元素数和所需的相等元素数。 输入的第二行包含n个整数a1，a2，…，an（1≤ai≤109），其中ai是a的第i个元素。

Output

打印一个整数-在数组中获得至少k个相等元素所需的最小移动次数。

题解

针对其中的某一个元素(大小i,出现次数为j)，如果我们想要将他的数量达到到K，共有四种情况

• 他本身的出现次数超过K次
• 将比它小的所有数字补充到i - 1,然后选取K - j个数字增大到i
• 将比它大的所有数字补充到i + 1,然后选取K - j个数字减小到i
• 将比它小的所有数字补充到i + 1,比它小的所有数字补充到i - 1,然后选取K - j个数字到i

除开情况1之外，其他情况的第一步操作（将其余数字补充到i±1）都是一个常数，可以直接用前缀和求出，求出之后O(n)的时间遍历并针对上面四种情况分别求出他们的花费取最小即可

代码

#include<bits/stdc++.h>
using namespace std;
#define PII pair<int,int>
#define fi first
#define se second
#define mp make_pair
#define LL long long
#define PLL pair<LL,LL>
const int maxn = 2e5 + 10;
const int mod = 1e9 + 7; 
int N,M,S;
int b[maxn];
PLL a[maxn];
LL presum[maxn],erpsum[maxn];
LL prenum[maxn],erpnum[maxn];
struct Node {
    LL precost; //前缀入门门槛
    LL prenum;  //前缀数量
    LL erpcost; //后缀入门门槛
    LL erpnum;  //后缀数量 
}node[maxn];
int main(){
    scanf("%d%d",&N,&M);
    for(int i = 1; i <= N ; i ++) {
        scanf("%d",&b[i]);
    }
    sort(b + 1,b + 1 + N);
    int cnt = 0;
    for(int i = 1; i <= N ; i ++) {
        a[++cnt].fi = b[i];
        int j;
        for(j = i; j <= N && b[i] == b[j]; j ++) a[cnt].se++;
        i = j - 1; 
    }
    for(int i = 1; i <= cnt ; i ++) {
        node[i].prenum = prenum[i - 1];
        node[i].precost = (a[i].fi - 1) * prenum[i - 1] - presum[i - 1];
        prenum[i] = prenum[i - 1] + a[i].se;
        presum[i] = presum[i - 1] + a[i].se * a[i].fi;
    }
    for(int i = cnt; i >= 1; i --) {
        node[i].erpnum = erpnum[i + 1];
        node[i].erpcost = erpsum[i + 1] - erpnum[i + 1] * (a[i].fi + 1);
        erpnum[i] = erpnum[i + 1] + a[i].se;
        erpsum[i] = erpsum[i + 1] + a[i].se * a[i].fi;
    }
    LL ans = 1e18;
    for(int i = 1; i <= cnt; i ++) {
        if(a[i].se >= M) {
            ans = 0;
            break;
        }
        if(a[i].se + node[i].erpnum >= M) {
            ans = min(ans,M - a[i].se + node[i].erpcost);
        }
        if(a[i].se + node[i].prenum >= M) {
            ans = min(ans,M - a[i].se + node[i].precost);
        }
        if(a[i].se + node[i].prenum + node[i].erpnum >= M) {
            ans = min(ans,M - a[i].se + node[i].precost + node[i].erpcost);
        }
    }
    printf("%lld\n",ans);
    return 0;
}