Cod sursa(job #2142169)

Utilizator preda.andreiPreda Andrei preda.andrei Data 24 februarie 2018 19:58:43
Problema Arbore partial de cost minim Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 3.1 kb
#include <fstream>
#include <tuple>
#include <vector>

using namespace std;

template <class T, class Compare>
class MinHeap
{
public:
  MinHeap() : comp_(Compare()) {}

  void Push(const T &val);

  T Pop();

  int Size() const { return vec_.size(); }

private:
  static int Father(int node) { return ((node - 1) >> 1); }
  static int LeftSon(int node) { return (node << 1) + 1; }
  static int RightSon(int node) { return (node << 1) + 2; }

  bool IsSmaller(int a, int b) { return comp_(vec_[a], vec_[b]); }

  void Percolate(int node);

  void Sift(int node);

  void Remove(int node);

  vector<T> vec_;
  Compare comp_;
};

template <class T, class Compare>
void MinHeap<T, Compare>::Push(const T &val)
{
  vec_.push_back(val);
  Percolate(Size() - 1);
}

template <class T, class Compare>
T MinHeap<T, Compare>::Pop()
{
  if (Size() < 1) {
    return T();
  }

  auto top = vec_[0];
  Remove(0);
  return top;
}

template <class T, class Compare>
void MinHeap<T, Compare>::Percolate(int node)
{
  while (node > 0 && IsSmaller(node, Father(node))) {
    auto father = Father(node);
    swap(vec_[node], vec_[father]);
    node = father;
  }
}

template <class T, class Compare>
void MinHeap<T, Compare>::Sift(int node)
{
  if (LeftSon(node) >= Size()) {
    return;
  }

  auto smallest = LeftSon(node);
  auto right = RightSon(node);

  if (right < Size() && IsSmaller(right, smallest)) {
    smallest = right;
  }

  if (IsSmaller(smallest, node)) {
    swap(vec_[smallest], vec_[node]);
    Sift(smallest);
  }
}

template <class T, class Compare>
void MinHeap<T, Compare>::Remove(int node)
{
  if (Size() < 1) {
    return;
  }

  swap(vec_[node], vec_[Size() - 1]);
  vec_.pop_back();

  if (node > 0 && IsSmaller(node, Father(node))) {
    Percolate(node);
  } else {
    Sift(node);
  }
}

using EdgeList = vector<pair<int, int>>;
using Graph = vector<EdgeList>;

using Tuple = tuple<int, int, int>;

pair<int, EdgeList> FindMst(const Graph &g)
{
  int nodes = g.size();

  EdgeList mst(nodes - 1);
  int mst_index = 0;
  int mst_cost = 0;

  vector<bool> used(nodes);
  MinHeap<Tuple, less<Tuple>> heap;

  heap.Push(make_tuple(0, 0, -1));

  while (heap.Size() > 0 && mst_index + 1 < nodes) {
    int cost, node, prev;
    tie(cost, node, prev) = heap.Pop();

    if (used[node]) {
      continue;
    }

    used[node] = true;
    mst_cost += cost;

    if (prev >= 0) {
      mst[mst_index++] = {node, prev};
    }

    for (const auto &p : g[node]) {
      int next, edge_cost;
      tie(next, edge_cost) = p;

      if (!used[next]) {
        heap.Push(make_tuple(edge_cost, next, node));
      }
    }
  }
  return {mst_cost, mst};
}

int main()
{
  ifstream fin("apm.in");
  ofstream fout("apm.out");

  int nodes, edges;
  fin >> nodes >> edges;

  Graph graph(nodes);
  for (int i = 0; i < edges; ++i) {
    int a, b, cost;
    fin >> a >> b >> cost;
    graph[a - 1].push_back({b - 1, cost});
    graph[b - 1].push_back({a - 1, cost});
  }

  auto mst_pair = FindMst(graph);
  fout << mst_pair.first << "\n";
  fout << nodes - 1 << "\n";

  for (const auto &e : mst_pair.second) {
    fout << e.first + 1 << " " << e.second + 1 << "\n";
  }
  return 0;
}