1014 Computation of Minimum Length of Pipeline

http://rose.u-aizu.ac.jp/onlinejudge/ProblemSet/description.jsp?id=1014

概要

温泉とそれぞれの地区の距離が与えられるので、すべての地区に温泉がいきわたるための
最小のパイプの長さを求める。

方針

はじめにヒープに温泉から出ている辺をすべて追加しておいてヒープから辺を取り出しながら、
もしすでに両方の辺に温泉がいきわたっているなら無視、片方だけならまだ温泉が届いていなかった側から出る辺をヒープに追加するということを繰り返す。

ソース

#include<iostream>
#include<vector>
#include<queue>
#include<cassert>
using namespace std;

#define REP(i,a,b) for(i=a; i<b; ++i)
#define rep(i,n) REP(i,0,n)

struct Edge{
  int s,t,c;
  Edge(int s,int t,int c) : s(s), t(t), c(c) {;}
};

struct UnionFind {
  vector<int> data;
  UnionFind(int size) : data(size, -1) { }
  bool unionSet(int x,int y) {
    x = root(x);
    y = root(y);
    if(x != y) {
      if(data[y] < data[x]) swap(x, y);
      data[x] += data[y];
      data[y] = x;
    }
      return x != y;
  }
  bool findSet(int x,int y) {
    return root(x) == root(y);
  }
  int root(int x) {
    return data[x] < 0 ? x : data[x] = root(data[x]);
  }
  int size(int x) {
    return -data[root(x)];
  }
};

bool operator<(const Edge &e,const Edge &f) {
  return (e.c != f.c) ? e.c > f.c : (e.s != f.s) ? e.s > f.c : e.t > f.t;
}

int solve(vector<vector<Edge> > &g,int hots) {
  const int n = g.size();
  int i,j;
  vector<int> atr(n,0);
  rep(i,hots) atr[i] = 1;
  UnionFind uf(n);

  priority_queue<Edge> Q;
  rep(i,hots) rep(j,g[i].size()) Q.push(g[i][j]);

  int ans = 0;
  while(!Q.empty()) {
    Edge e = Q.top(); Q.pop();
    if(uf.findSet(e.s,e.t)) continue;
    if(atr[e.s] != 0 && atr[e.t] != 0) continue;
    ans += e.c;
    //cout<<"add "<<e.s<<" "<<e.t<<endl;
    uf.unionSet(e.s,e.t);
    if(atr[e.s] == 0 && atr[e.t] == 0) assert(false);
    if(atr[e.s] == 0) {
      atr[e.s] = 1;
      rep(i,g[e.s].size()) Q.push(g[e.s][i]);
    }else if(atr[e.t] == 0) {
      atr[e.t] = 1;
      rep(i,g[e.t].size()) Q.push(g[e.t][i]);
    }
  }
  return ans;
}

int main() {
  int s,d,i,j,k;
  while(cin>>s>>d, s|d) {
    vector<vector<Edge> > g(s+d);
    rep(i,s) {
      rep(j,d) {
	cin>>k;
	if(!k) continue;
	g[i].push_back(Edge(i,j+s,k));
      }
    }

    rep(i,d) {
      REP(j,i+1,d) {
	cin>>k;
	if(!k) continue;
	g[i+s].push_back(Edge(i+s,j+s,k));
	g[j+s].push_back(Edge(j+s,i+s,k));
      }
    }

    cout<<solve(g,s)<<endl;
  }
}