Codeforces 371E Subway Innovation (前缀和预处理应用)

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b) for(int i(a); i <= (b); ++i)
#define LL           long long

const int N = 300010;

struct node{
    LL x, y;
    friend bool operator < (const node &a, const node &b){
        return a.x < b.x;
    }
} a[N];

LL x[N], s[N], delta[N];
LL n, k, p, ss, ans, cnt;

int main(){

    scanf("%lld", &n);
    s[0] = 0;
    rep(i, 1, n){
        scanf("%lld", &a[i].x);
        a[i].y = i;
    }
    
    scanf("%lld", &k);
    sort(a + 1, a + n + 1);
    rep(i, 1, n) x[i] = a[i].x;
    rep(i, 1, n) s[i] = s[i - 1] + x[i];
    int j = 0;
    rep(i, k + 1, n){
        ++j;
        delta[i] = (k - 1) * (x[i] + x[j]) - 2 * (s[i - 1] - s[j]);
    }
    cnt = ans = 0, p = k;
    rep(i, k + 1, n){
        cnt += delta[i];
        if (ans > cnt){
            ans = cnt;
            p = i;
        }
    }
    rep(i, p - k + 1, p) printf("%lld\n", a[i].y);
    return 0;

}