yvfRP6 - Online C++0x Compiler & Debugging Tool

#include<bits/stdc++.h>
#define ii pair<int, int>
#define fi first
#define se second
#define iii pair<ii, int>
using namespace std;

const int MOD = int(1e5 + 7);
const int nmax = int(1e6 + 7);

int n, m;
vector<int> a;
long long d[nmax] = {0};

void prepare() {
}
void solve() {
}
int find_first(int u, int v, int x) {
    int L = u, R = v;
    while(true) {
        if(L == R) break;
        if(L == R-1) {
            if(a[L] < x) L = R;
            break;
        }
        int m = (L + R) / 2;
        if(a[m] >= x) R = m;
        else L = m+1;
    }
    return L;
}
void ip() {
    cin >> n >> m;
    a.resize(n+1);
    for(int i = 1; i <= n; ++i) {
        cin >> a[i];
    }
    sort(a.begin()+1, a.end());
    for(int i = 1; i <= n; ++i) {
        d[i] = d[i-1] + a[i];
    }
    long long ans = LLONG_MAX;
    for(int i = 1; i <= n - m + 1; ++i) {
//       select from i -> i + m - 1
        int tbc = (d[i + m -1] - d[i-1]) / m;
        int pos = find_first(i, i + m -1, tbc);
        long long thieu = 0, thua = 0;
        if(pos > i) {
            thieu = 1LL * (pos -1 -i +1) * tbc - (d[pos-1] - d[i-1]);
        }
        if(pos <= i + m -1) {
            thua = abs(1LL * (i + m -1 -pos +1) * tbc - (d[i + m -1] - d[pos-1]));
        }
        if(1LL * tbc * m <= d[n])
            ans = min(ans, max(thua, thieu));
        tbc = (d[i + m -1] - d[i-1]) / m + 1;
        pos = find_first(i, i + m -1, tbc);
        thieu = 0, thua = 0;
        if(pos > i) {
            thieu = 1LL * (pos -1 -i +1) * tbc - (d[pos-1] - d[i-1]);
        }
        if(pos <= i + m -1) {
            thua = abs(1LL * (i + m -1 -pos +1) * tbc - (d[i + m -1] - d[pos-1]));
        }
        if(1LL * tbc * m <= d[n])
            ans = min(ans, max(thua, thieu));
//        for(int j = i; j <= i + m - 1; ++j) {
//            cout << a[j] << " ";
//        } cout << '\n';
//        cout << "TBC: " << tbc << '\n';
//        cout << thua << " | " << thieu << '\n';
    }
    cout << ans;
}
int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);

    if(fopen("task.inp", "r")) {
        freopen("task.inp", "r", stdin);
        freopen("task.out", "w", stdout);
    }
    ip();
}