集合操作

参考了大佬的代码:boringhaker 提交的代码

先將序列a排序。

如果 a[pos] 后边的数减去 k * p的平均值大于等于 a[pos - 1],就不需要管pos前边的数。
用后缀和找到 pos 的位置。

pos后边的数减去 (a[i] - key)/p*p
如果 k 依旧不为0,继续从大到小减去p。

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N = 1e6+10;
int a[N];

signed main(){
    std::ios_base::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    int n, k, p;
    cin>>n>>k>>p;
    for(int i = 1; i <= n; i++){
        cin>>a[i];
    }
    sort(a+1, a+1+n);

    // 特判 
    if(p == 0){
        for(int i = 1; i <= n; i++){
            cout<<a[i]<<" ";
        }
        return 0;
    }

    int pos, key, sum = 0;
    a[0] = -LLONG_MAX;

    // 找到pos的位置 
    for(int i = n; i >= 1; i--){
        sum += a[i];
        if((sum - k*p)/(n-i+1) >= a[i-1]){
            pos = i;
            key = (sum - k*p)/(n-i+1);
            break;
        }
    }

    for(int i = n; i >= pos; i--){
        k -= (a[i] - key)/p;
        a[i] -= (a[i] - key)/p*p;
    }
    sort(a+1, a+1+n);

    for(int i = n; i >= pos; i--){
        if(k > 0){
            a[i] -= p;
            k--;
        }
        else{
            break;
        }
    }
    sort(a+1, a+1+n);

    for(int i = 1; i <= n; i++){
        cout<<a[i]<<" ";
    }

    return 0;
}