Cod sursa(job #2206788)

Utilizator georgeromanGeorge Roman georgeroman Data 23 mai 2018 20:14:04
Problema Arbore partial de cost minim Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 1.83 kb
#include <algorithm>
#include <fstream>
#include <set>
#include <vector>

class WeightedUF {
  std::vector<int> id;
  std::vector<int> size;
  int count;
  
public:
  WeightedUF(int n) {
    id.reserve(n);
    size.reserve(n);
    count = n;
    for (int i = 0; i < n; i++) {
      id.push_back(i);
      size.push_back(1);
    }
  }
  
  int find(int p) {
    while (p != id[p]) p = id[p];
    return p;
  }
  
  void dounion(int p, int q) {
    int pId = find(p);
    int qId = find(q);
    if (pId == qId) return;
    if (size[pId] < size[qId]) {
      id[pId] = qId;
      size[qId] += size[pId];
    } else {
      id[qId] = pId;
      size[pId] += size[qId];
    }
    count--;
  }
  
  bool connected(int p, int q) {
    return find(p) == find(q);
  }
  
  int getCount() {
    return count;
  }
};

struct Edge {
  int u, v, cost;
  Edge(int u, int v, int cost): u(u), v(v), cost(cost) {}
};

int main() {
  std::ifstream in("apm.in");
  std::ofstream out("apm.out");
  
  int n, m;
  in >> n >> m;
  std::vector<Edge> edges;
  edges.reserve(m);
  for (int i = 0; i < m; i++) {
    int u, v, cost;
    in >> u >> v >> cost;
    edges.push_back(Edge(u, v, cost));
  }
  
  std::vector<Edge> minTree;
  std::vector<std::set<int>> whichTree;
  for (int i = 0; i < n; i++) {
    whichTree.push_back(std::set<int>());
  }
  
  std::sort(edges.begin(), edges.end(), [](Edge& e1, Edge& e2){ return e1.cost > e2.cost; });
  WeightedUF set(n);
  int totalCost = 0;
  while (set.getCount() != 1) {
    Edge& best = edges[edges.size() - 1];
    edges.pop_back();
    if (!set.connected(best.u - 1, best.v - 1)) {
      set.dounion(best.u - 1, best.v - 1);
      totalCost += best.cost;
      minTree.push_back(best);
    }
  }
  
  out << totalCost << "\n";
  out << minTree.size() << "\n";
  for (Edge& e : minTree) {
    out << e.u << " " << e.v << "\n";
  }
  
  return 0;
}