Codeforces 371(#218 Div2 only) E. Subway Innovation

問題

数直線上に駅がn個ある。i番目の駅はx[i]にある。
ここから駅をk個残してあとは廃止する。

残す駅は、全ての二つの駅間の距離の平均が最小になるようにしたい。
（Sを残す駅としたら、Σ[a, b∈S, a < b]|x[a] - x[b]| / (k * (k - 1) / 2)を最小にする）

駅の残し方をどれか一通り出力せよ。

制約条件

n≦3*10^5
-10^8≦x[i]≦10^8
2≦k≦n-1

方針

まず、個数が一定なので合計を最小化するとしてよい。
そして、残す駅は明らかに連続したk個。

なので始点をn-k+1通り試す。
始点を一個ずらしたときの距離の合計の変化をすばやく計算できたらよい。

x[i]を入れて、x[i - k]を外すと考えると、
外すことによって、
(x[i - k + 1] - x[i - k]) + (x[i - k + 2] - x[i - k]) + …が減り、
入れることによって、
(x[i] - x[i - k + 1]) + (x[i] - x[i - k + 2]) + ……が増える。

累積和を持っておくことで、O(1)でこの差分を計算することができる。

ソースコード

const int N = 300010;

int n, k, pos[N];
ll sum[N];

vector<pi> v;

int main(){
	
	cin >> n;
	rep(i, n) cin >> pos[i], v.pb(mp(pos[i], i + 1));
	cin >> k;
	
	sort(all(v));
	
	rep(i, n){
		pos[i] = v[i].first;
		sum[i + 1] = sum[i] + pos[i];
	}
	
	ll tmp = 0;
	
	rep(i, k) tmp += (ll)i * pos[i] - sum[i];
	
	ll ans = tmp;
	int ansi = 0;
	
	for(int i = k; i < n; i++){
		
		tmp -= sum[i] - sum[i - k + 1] - (k - 1ll) * pos[i - k];
		tmp += (k - 1ll) * pos[i] - sum[i] + sum[i - k + 1];
		
		if(ans > tmp) ans = tmp, ansi = i - k + 1;
	}
	
	rep(i, k) cout << v[ansi + i].second << (i==k-1?"\n":" ");
	
	return 0;
}