Minimum Spanning Tree

给十几个城市供电，连接不同城市的花费不同，让花费最小同时连到所有的边。给出一系列connection类，里面是edge两端的城市名和它们之间的一个cost，找出要你挑一些边，把所有城市连接起来并且总花费最小。不能有环，最后所以城市要连成一个连通块。
不能的话输出空表，最后还要按城市名字排序输出，按照node1来排序,如果一样的话再排node2。
输入:

{"Acity","Bcity",1}
("Acity","Ccity",2}
("Bcity","Ccity",3}

输出：

("Acity","Bcity",1}
("Acity","Ccity",2}

补充一句，test case一共有6个。

Solution

#include <iostream>
#include <string>
#include <vector>
#include <queue>
#include <unordered_map>
#include <unordered_set>
using namespace std;

struct connection {
  int start, end;
  int weight;
  connection(int i, int j, int k): start(i), end(j), weight(k) {}
};

bool compare(connection* a, connection* b) {
  return a->weight < b->weight;
}
class Solution {
  public:
    vector<connection*> mst(vector<connection*>& edges) {
      sort(edges.begin(), edges.end(), compare);
      int n = edges.size();
      unordered_set<int> included;
      unordered_map<int, int> setNumber;
      unordered_map<int, unordered_set<int> > sets;
      vector<connection*> res;
      for ( int i = 0; i < n; i++ ) {
        int j = edges[i]->start, k = edges[i]->end;
        if ( included.find(j) != included.end() and included.find(k) != included.end() ) {
          if ( setNumber[j] != setNumber[k] ) {
            res.push_back(edges[i]);
            // merge k to j
            int jj = setNumber[j], kk = setNumber[k];
            for ( int ii: sets[kk] ) setNumber[ii] = jj;
            sets.erase(sets.find(kk));
          }
          continue;
        }
        res.push_back(edges[i]);
        if ( included.find(j) == included.end() and included.find(k) == included.end() ) {
          sets[j].insert(j);
          sets[j].insert(k);
          setNumber[j] = j; setNumber[k] = j;
        }
        else {
          if ( included.find(j) != included.end() ) {
            sets[setNumber[j]].insert(k);
            setNumber[k] = setNumber[j];
          }
          else {
            sets[setNumber[k]].insert(j);
            setNumber[j] = setNumber[k];
          }
        }
        included.insert(j); included.insert(k);
      }
      if ( sets.size() != 1 ) {
        res.clear();
        return res;
      }
      for ( unordered_set<int>::iterator it = included.begin(); it != included.end(); it++ ) {
        if ( setNumber[*it] != -1 ) continue;
        res.clear();
        return res;
      }
      return res;
    }
};
int main() {
  Solution sol;
  vector<connection*> edges;

  edges.push_back(new connection(1, 2, 6));
  edges.push_back(new connection(2, 3, 4));
  edges.push_back(new connection(3, 4, 5));
  edges.push_back(new connection(4, 5, 8));
  edges.push_back(new connection(5, 6, 2));
  edges.push_back(new connection(2, 6, 10));
  edges.push_back(new connection(5, 3, 9));
  edges.push_back(new connection(6, 3, 7));
  edges.push_back(new connection(2, 5, 3));
  edges.push_back(new connection(1, 6, 16));
  vector<connection*> res = sol.mst(edges);
  for ( int i = 0; i < int(res.size()); i++ ) {
    cout << res[i]->start << " to " << res[i]->end << ": " << res[i]->weight << endl;
  }
  return 0;
}