Codeforces 1328F Make k Equal

Description

给你一个数组及其长度 $n$，以及一个值 $k$。

仅通过将「最小元素自增 $1$」或「最大元素自减 $1$」的方式，以使得至少有 $k$ 个元素相等，输出所需的最少步数。

Solution

如果一个数组中已经有一种数 $\ge k$ 个了，那么答案一定是 $0$。

否则的话，显然我们最终要凑的数字一定是数组当中的一个数。

所以我们就可以枚举这个数。

数组先排个序，分三种情况，一种是只动左边，一种是只动右边，还有一种是两边都动。

如果只动前面，把 $a_1 \sim a_i$ 全部弄成 $a_i$ 需要多少步？

答案是 $i \cdot a_i- \sum\limits_{j=1}^{i}a_j$。

但是这样就一共有了 $i$ 个 $a_i$。如果 $i = k$，那么正好；如果 $i < k$，那么是不符合条件的；如果 $i > k$，可以发现，我们又浪费了一部分，即多了 $i - k$ 个，那么怎么办呢？很简单，让这 $i - k$ 个保持在 $a_{i} - 1$ 就好了，所以用：

$$i \cdot a_i- \left(\sum_{j=1}^{i}a_j\right) - (i - k)$$

来更新答案。

只动后面同理，条件是 $n - i + 1 \ge k$，用：

$$\left(\sum_{j=i}^{n}a_j \right)- (n - i + 1) \cdot a_i - \left((n - i + 1) - k\right)$$

来更新答案。

如果同时 $i<k$ 且 $n-i+1<k$，那么两边就都要动了，这时相等的数的个数恰好为 $n$，很显然，把上面两个式子拼起来即可，用：

$$i \cdot a_i- \left(\sum_{j=1}^{i}a_j\right) +\left(\sum_{j=i}^{n}a_j \right)- (n - i + 1) \cdot a_i - (n-k)$$

来更新答案。

$\sum\limits_{j=1}^{i}a_j$ 和 $\sum\limits_{j=i}^{n}a_j$ 可以预处理，所以时间复杂度是 $\mathcal O(n)$。

代码供参考。

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 2e5 + 5;
int n, k, a[N], pre[N], suf[N], ans = LLONG_MAX, cnt[N];
signed main()
{
    cin >> n >> k;
    for(int i = 1; i <= n; i++) cin >> a[i];
    sort(a + 1, a + n + 1);
    for(int i = 1; i <= n; i++) pre[i] = pre[i - 1] + a[i];
    for(int i = n; i >= 1; i--) suf[i] = suf[i + 1] + a[i];
    for(int i = 1; i <= n; i++)
    {
        if(a[i] == a[i - 1]) cnt[i] = cnt[i - 1] + 1;
        else cnt[i] = 1;
        if(cnt[i] >= k) return puts("0") && 0;
    } 
    for(int i = 1; i <= n; i++)
    {
        if(i >= k) ans = min(ans, i * a[i] - pre[i] - (i - k));
        if(n - i + 1 >= k) ans = min(ans, suf[i] - (n - i + 1) * a[i] - (n - i + 1 - k));
        ans = min(ans, i * a[i] - pre[i] + suf[i] - (n - i + 1) * a[i] - (n - k));
    } 
    cout << ans << endl;
    return 0;
}