CodeChef-RHOUSE : Houses and Restaurants

問題概要

N(<10^5)個の建物がある。建物は家かレストランのいずれかである。辺の候補がM(<4*10^5)本ある。辺は重み付きの無向辺で、重みは負もとりうる。全ての家がいずれかのレストランとつながるように辺を選んでコストを最小化する問題。

解法

一度レストランとつながった家はレストランになると考えれば分かりやすい。コストを昇順に持つプライオリティーキューを用意して、片方がレストランになってるような辺を随時加えていく。この際家がレストランになったならその家を端点に持つ家同士をつなぐ辺をプライオリティーキューに追加する。最後にコストが負の未使用な辺を全部集めればOK。
なお、コストが負の辺を最初に選ぶ(確定させる)のはかえってやりにくくなる。

acceptされたコード

計算量O(M * log M)。

#include <cstdio>
#include <queue>
#include <algorithm>
using namespace std;

typedef long long int64;

const int MAX_N = 1e5;
const int MAX_M = 4e5;

int N, M, R;
char which[MAX_N + 1];

struct edge{
	int from, to, cost;
	bool used;
};

struct id{
	int idx;
};

edge es[MAX_M];

bool operator<(const id& lhs, const id& rhs){
	return es[lhs.idx].cost > es[rhs.idx].cost;
}

struct UnionFind{
	static const int MAX_UF = MAX_N + 1;

	int pars[MAX_UF];
	int ranks[MAX_UF];

	void init(int n){
		for(int i=0; i<n; i++){
			pars[i] = i;
			ranks[i] = 0;
		}
	}

	int get_root(int x){
		if(pars[x] == x){
			return x;
		}
		else{
			return pars[x] = get_root(pars[x]);
		}
	}

	void merge(int x, int y){
		x = get_root(x);
		y = get_root(y);
		if( x == y ){
			return ;
		}

		if(ranks[x] < ranks[y]){
			pars[x] = y;
		}
		else if(ranks[y] < ranks[x]){
			pars[y] = x;
		}
		else{
			pars[x] = y;
			ranks[y]++;
		}
	}

	bool is_same(int x, int y){
		return get_root(x) == get_root(y);
	}
};

UnionFind uf;

vector<id> ids[MAX_N];

int64 solve(){
	int64 ans = 0;
	scanf("%d%d ", &N, &M);
	scanf("%[^\n]%*c", which);
	R = N;
	uf.init( N + 1 );
	for(int i=0; i<N; i++){
		if(which[i] == 'R'){
			uf.merge(i, R);
		}
	}
	for(int i=0; i<N; i++){
		ids[i].clear();
	}


	for(int i=0; i<M; i++){
		int a, b, c;
		scanf("%d%d%d", &a, &b, &c);
		a--; b--;
		es[i] = (edge){a, b, c, false};
	}

	priority_queue<id> que;

	for(int i=0; i<M; i++){
		int a = es[i].from, b = es[i].to, cost = es[i].cost;

		if(uf.is_same(R, a) ^ uf.is_same(R, b)){
			que.push((id){i});
		}
		else if(!uf.is_same(R, a)){
			ids[a].push_back((id){i});
			ids[b].push_back((id){i});
		}

	}

	while(!que.empty()){
		edge& e = es[que.top().idx];
		que.pop();
		int a = e.from, b = e.to, cost = e.cost;
		if( uf.is_same(R, a) && uf.is_same(R, b) ){
			continue;
		}
		//aが新規参入になるように調整
		if(uf.is_same(R, a)){
			swap(a, b);
		}
		uf.merge(a, b);
		ans += cost;
		e.used = true;

		for(int i=0; i<(int)ids[a].size(); i++){
			edge& e = es[ids[a][i].idx];
			int tar = e.from;
			if(tar == a){
				tar = e.to;
			}
			if( !uf.is_same(R, tar) ){
				que.push(ids[a][i]);
			}
		}
	}

	for(int i=0; i<M; i++){
		if(es[i].cost < 0 && !es[i].used){
			ans += es[i].cost;
		}
	}

	return ans;
}

int main(){
	int T;
	scanf("%d", &T);
	while(T--){
		printf("%lld\n", solve());
	}

	return 0;
}