SRM 492 Div2 Hard TimeTravellingSalesman

問題

http://apps.topcoder.com/stat?c=problem_statement&pm=11049&rd=14245

N個の街があり、セールスマンが街0にいるとする。N個の街の間の道の情報がvector string で与えられる。これを一つのstringに結合した時に"a0,b0,cost0 a1,b1,cost1 ......"という形を取り、街aiとbiをcostiの距離の道でつながれていることを示す。この道は双方向である。セールスマンはタイムマシンを使用して今までに訪れた街に訪れた順と逆順に瞬時に移動できる。つまりA->B->C->D->Eと訪れた時、Cに瞬時に移動できる。またそこからB、Aに移動できるがD、Eには移動できない。セールスマンがすべての街に移動するのに必要な最小移動距離を求めよ。

やりかた

移動経路は木になるようにできるので、最小全域木にすることで最小移動距離にできる。

以下ソース。最小全域木はクラスカル法。

struct edge {
  int from/*自分*/, to/*行き先*/, flow/*流量*/; ll cost/*費用*/;
  edge(int from, int to, int flow, ll cost) :
    from(from), to(to), flow(flow), cost(cost) {}
  edge(int to, ll cost) :
    to(to), cost(cost) {}
  bool operator >(const edge &e) const{
    return cost > e.cost;
  }
  bool operator <(const edge &e) const{
    return cost < e.cost;
  }
};

typedef vector<vector<edge> > Graph;
typedef pair<int, int> P;

int par[10000], rrank[10000];

void uf_init(int n){
  for(int i = 0; i < n; i++){
    par[i] = i;
    rrank[i] = 0;
  }
}

int uf_find(int x){
  if(par[x] == x){
    return x;
  }else{
    return par[x] = uf_find(par[x]);
  }
}

void uf_unite(int x, int y){
  x = uf_find(x);
  y = uf_find(y);
  if(x == y) return;

  if(rrank[x] < rrank[y]){
    par[x] = y;
  }else{
    par[y] = x;
    if(rrank[x] == rrank[y]) rrank[x]++;
  }
}

bool uf_same(int x, int y){
  return uf_find(x) == uf_find(y);
}
// クラスカル法
pair<ll, vector<edge> > kruskal(const Graph &g){
  int V = g.size();
  uf_init(V);
  priority_queue<edge, vector<edge>, greater<edge> > Q;
  for(int i = 0; i < V; i++) 
    for(int j = 0; j < (int)g[i].size(); j++)
      if(i < g[i][j].to) Q.push(g[i][j]);

  ll total = 0;
  vector<edge> MST;
  while((int)MST.size() < V - 1 && !Q.empty()){
    edge e = Q.top(); Q.pop();
    if(!uf_same(e.from, e.to)){
      uf_unite(e.from, e.to);
      MST.push_back(e);
      total += e.cost;
    }
  }
  return pair<ll, vector<edge> >(total, MST);
}


class TimeTravellingSalesman {
public:
  vector<int> parse(vector<string> s){
    string t;
    for(int i = 0; i < (int)s.size(); i++){
      for(int j = 0; j < (int)s[i].length(); j++){
	if(isdigit(s[i][j])) t += s[i][j];
	else t += " ";
      }
    }
    stringstream ss;
    ss << t;
    int u;
    vector<int> ret;
    while(ss >> u) ret.push_back(u);
    return ret;
  }
  long long determineCost(int N, vector <string> roads) {
    vector<int> val = parse(roads);
    cout << val.size() << endl;
    Graph g;
    g.resize(N);
    for(int i = 0; i < N; i++)
      for(int j = 0; j < N; j++)
	g[i].push_back(edge(i, j, 0, INF));

    for(int i = 0; i < (int)val.size(); i += 3){
      g[val[i]][val[i + 1]].cost = val[i + 2];
      g[val[i + 1]][val[i]].cost = val[i + 2];
    }

    pair<ll, vector<edge> > MSP = kruskal(g);
    return MSP.first >= INF ? -1 : MSP.first;
  }
};

変な入力方法だが、SRM上では何回か見ているのでライブラリにしておこう。
あとDiv2 Easyがsystestに落ちていたので慌てて修正。

追記（2014/02/24）
プリム法でも実装

int s2i(string s){
  stringstream ss;
  ss << s;
  int a;
  ss >> a;
  return a;
}

ll prim(vector<vector<ll> > dist){
  int N = dist.size();
  vector<bool> used(N, false);
  ll ret = 0;
  priority_queue<pair<ll, ll>, vector<pair<ll, ll> >, greater<pair<ll, ll> > > Q;
  Q.push(make_pair(0, 0));
  while(!Q.empty()){
    ll cost = Q.top().first;
    ll to = Q.top().second;
    Q.pop();
    if(used[to]) continue;
    used[to] = true;
    ret += cost;
    for(int i = 0; i < N; i++) Q.push(make_pair(dist[to][i], i));
  }
  return ret;
}

class TimeTravellingSalesman {
public:
  long long determineCost(int N, vector <string> roads) {
    string str = accumulate(ALL(roads), string(""));
    vector<vector<ll> > d(N, vector<ll>(N, INF));
    for(int i = 0; i < (int)str.length(); i++) if(str[i] == ',') str[i] = ' ';
    
    stringstream ss; ss << str;
    string s1, s2, s3;
    while(ss >> s1 >> s2 >> s3)
      d[s2i(s1)][s2i(s2)] = d[s2i(s2)][s2i(s1)] = s2i(s3);
    
    ll result = prim(d);
    return result >= INF ? -1 : result;
  }
};