Codeforces #647 Div1 D. Johnny and James

もう少し問題文短くなってくれるといいんだけどな。
http://codeforces.com/contest/1361/problem/D

問題

2次元座標上にN個の基地がある。また、座標の原点には中央基地がある。
基地同士は以下の通信方法がとれる。

• 直接通信ができるのは、原点方向または逆方向にある基地のみ。
• 中央基地を介して、間接的に通信することもできる。

N個の基地のうちK個だけ残したとき、基地の対における通信距離の総和を求めよ。
(なお、中央基地は残さないこともできる。ただしその場合も間接的な通信には関われる)

解法

まず基地を原点からの偏角毎にグループ化し、原点からの距離でソートしておく。
各グループ毎に見ていくと、距離を最大化するなら残す基地は原点寄りと遠方寄りから順に選ぶのが最適である。
そこで、尺取り法の要領で「合計で何個残すと、距離の合計の最大値はどうなるか」を求めておこう。

次に、その差分を取ると残す基地を1つ変えたとき、どの程度距離の合計を増やせるかがわかる。
全グループにおいて、その値をpriority_queueに入れ、大きい順に取り出そう。

int N,K;
map<pair<int,int>,vector<double>> M;
vector<vector<double>> V;
int id[501010];
double ret;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>x>>y;
		double r=hypot(x,y);
		if(r==0) continue;
		int g=__gcd(abs(x),abs(y));
		x/=g;
		y/=g;
		M[{x,y}].push_back(r);
	}
	
	FORR(m,M) {
		sort(ALL(m.second));
		vector<double> W;
		W.push_back(0);
		int L=0,R=0;
		double LS=0,RS=0,sum=0;;
		while(L+R<min(K,(int)m.second.size())) {
			double ml=m.second[L];
			double mr=m.second[m.second.size()-1-R];
			double ladd=(LS+ml+RS)*(K-(L+R+1))+(ml*L-LS)+(RS-ml*R)+sum;
			double radd=(LS+mr+RS)*(K-(L+R+1))+(mr*L-LS)+(RS-mr*R)+sum;
			//cout<<L<<" "<<R<<" "<<ladd<<" "<<radd<<endl;
			if(ladd>radd) {
				LS+=ml;
				L++;
				W.push_back(ladd);
				sum+=(ml*L-LS)+(RS-ml*R);
			}
			else {
				RS+=mr;
				R++;
				W.push_back(radd);
				sum+=(mr*L-LS)+(RS-mr*R);
			}
		}
		
		V.push_back(W);
	}
	V.push_back({0,0});
	N=V.size();
	priority_queue<pair<double,int>> Q;
	FOR(i,N) {
		Q.push({V[i][1]-V[i][0],i});
	}
	while(K--) {
		ret+=Q.top().first;
		x=Q.top().second;
		//cout<<"!"<<Q.top().first<<" "<<x<<endl;
		Q.pop();
		id[x]++;
		if(id[x]+1<V[x].size()) {
			Q.push({V[x][id[x]+1]-V[x][id[x]],x});
		}
		
	}
	_P("%.12lf\n",ret);
	
}

まとめ

やることはわかっても若干面倒くさい問題。