Codeforces #485 A. Fair

なんとか赤文字に戻れた。
http://codeforces.com/contest/986/problem/A

問題

N個の頂点の連結無向グラフが与えられる。
各頂点には1～Kの値が振られている。

各頂点において、1～Kの値のうち異なるS個の値を取るようなS個の頂点を選び、各頂点までの最短距離の和を最小化したい。
それぞれの最小値を求めよ。

解法

Kが小さいので、i=1～Kに対し、値がiであるような頂点群を始点とした各頂点までの最短距離をダイクストラ法で求めよう。
辺の距離は1固定なので、priority_queueは不要で高速に処理できる。

これらの結果を合わせると、各頂点において最寄りの1～Kの値を持つ頂点までの距離がわかるので、小さい順S個の和を返せばよい。

int N,M,K,S;
vector<int> E[101010];
int A[101010];
int D[101010][101];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K>>S;
	FOR(i,N) cin>>A[i],A[i]--;
	FOR(i,M) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,K) {
		queue<int> Q;
		FOR(x,N) {
			if(A[x]==i) Q.push(x);
			else D[x][i]=1<<20;
		}
		while(Q.size()) {
			x=Q.front();
			Q.pop();
			FORR(e,E[x]) if(D[e][i]>D[x][i]+1) {
				D[e][i]=D[x][i]+1;
				Q.push(e);
			}
		}
	}
	FOR(i,N) {
		vector<int> V;
		FOR(j,K) V.push_back(D[i][j]);
		sort(ALL(V));
		ll ret=0;
		FOR(j,S) ret+=V[j];
		cout<<ret<<" ";
	}
	cout<<endl;
	
}

まとめ

タイプミスで無駄に1WAした…。