Codeforces #218 Div2. E. Subway Innovation

pretestで1ミスしたけど、修正してなんとかAC。
http://codeforces.com/contest/371/problem/E

問題

1次元の路線上にN個の駅があり、それぞれの位置が与えられる。
ここで、N個のうちK個の駅を残し、K個の駅のうち任意の2駅を選んだときの平均距離が最短になるようにしたい。
残す駅K個を答えよ。

解法

まず駅を位置順でソートする。
あえて遠い駅同士をK個に残す必要はないので、K個残す候補は駅を端から1～N番と番号をふった場合、(1～K)、(2～(K+1))、(3～(K+2))、…、((N-K+1)～N）番の残し方のいずれかである。
あとは各区間内の合計距離が最短のものを残せばよい。

事前に各駅の座標の部分和を取っておけば、最初の(1～K)の合計距離はO(K)かかるけど、以降1つずつ区間をずらした合計距離はO(1)で求まるので全体でO(N)。
部分和計算も合計距離計算もO(N)なので結局O(N)で処理できる。

下記コードは平均距離でなく合計距離だけで処理を進めたところ、64bit整数でもあふれたのでdouble型で回避した。

int N,K;
ll X[400001],A[400001],S[400001];
pair<ll,int> ST[400001];
double ave[400001];

void solve() {
	int f,i,j,k,l,x,y;
	cin>>N;
	FOR(i,N) {
		cin>>X[i];
		X[i]+=100000000;
		ST[i]=make_pair(X[i],i+1);
	}
	sort(ST,ST+N);
	FOR(i,N) {
		A[i]=ST[i].first;
		S[i+1]=S[i]+A[i];
	}
	
	cin>>K;
	double tot=0;
	FOR(i,K) {
		tot += (S[K]-S[i+1])-(A[i]*(K-(i+1)));
		tot += (A[i]*(i))-S[i];
	}
	ave[0] = tot;
	for(i=K;i<N;i++) {
		tot -= 2*((S[i]-S[i-K+1])-(A[i-K]*(K-1)));
		tot += 2*((A[i]*(K-1))-(S[i]-S[i-K+1]));
		ave[i-(K-1)] = tot;
	}
	double mi=1e200;
	x=-1;
	FOR(i,N-K+1) if(ave[i]<mi) x=i,mi=ave[i];
	FOR(i,K) _P("%d ",ST[x+i].second);
	_P("\n");
}

まとめ

ここらへんも過去に覚えたテクニックがつかえて、成長が感じられる回でした。