Cod sursa(job #3146885)

Utilizator daristyleBejan Darius-Ramon daristyle Data 22 august 2023 23:54:50
Problema Arbore partial de cost minim Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.86 kb
#include <fstream>
#include <vector>

using namespace std;

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

struct Node{
		int node;
		short cost;
};
const int N_MAX = 2e5;
const int COST_MAX = 1e3;
const int HEAP_SIZE_MAX = N_MAX;
const int HEAP_VALUE_MAX = N_MAX;
vector<int> children[N_MAX];
vector<Node> edge[N_MAX];
bool inHeap[N_MAX];
int apm_cost = 0;

class OptimizedHeap{
private:
		Node heap[HEAP_SIZE_MAX];
		int pos[HEAP_VALUE_MAX];
		int heapSize = 0;

		bool (* cmp)(Node, Node);

		inline int parent(int node){
			return (node - 1) / 2;
		}

		inline int leftChild(int node){
			return 2 * node + 1;
		}

		inline int rightChild(int node){
			return 2 * node + 2;
		}

		void upHeap(int node){
			while(node && cmp(heap[node], heap[parent(node)])){
				swap(heap[node], heap[parent(node)]);
				pos[heap[node].node] = node;
				pos[heap[parent(node)].node] = parent(node);

				node = parent(node);
			}
		}

		void downHeap(int node){
			int smallest = node, left = leftChild(node), right = rightChild(node);

			if(left < heapSize && cmp(heap[left], heap[smallest]))
				smallest = left;
			if(right < heapSize && cmp(heap[right], heap[smallest]))
				smallest = right;

			if(smallest != node){
				swap(heap[smallest], heap[node]);
				pos[heap[node].node] = node;
				pos[heap[smallest].node] = smallest;

				downHeap(smallest);
			}
		}

public:

		OptimizedHeap(bool (* _cmp)(Node, Node)) : cmp(_cmp){}

		void insert(Node x){
			heap[heapSize] = x;
			pos[x.node] = heapSize;
			++heapSize;

			upHeap(heapSize - 1);
		}

		Node query(){
			return heap[0];
		}

		void extractRoot(){
			swap(heap[0], heap[heapSize - 1]);
			pos[heap[0].node] = 0;
			--heapSize;
			downHeap(0);
		}

		void update(int node, short cost){
			if(cmp({node, cost}, heap[pos[node]]))
				heap[pos[node]].cost = cost;

			upHeap(pos[node]);
			downHeap(pos[node]);
		}

		bool empty(){
			return !heapSize;
		}
};

bool cmp(Node a, Node b){
	return a.cost < b.cost;
}

OptimizedHeap heap(cmp);

void Prim(int nodes){
	for(int i = 0; i < nodes; ++i){
		heap.insert({i, COST_MAX + 1});
		inHeap[i] = true;
	}

	heap.update(0, 0);

	for(int i = 0; i < nodes; ++i){
		Node current = heap.query();
		heap.extractRoot();

		inHeap[current.node] = false;
		apm_cost += current.cost;

		for(auto neighbour: edge[current.node])
			if(inHeap[neighbour.node])
				heap.update(neighbour.node, neighbour.cost);
			else if(neighbour.cost == current.cost)
				children[neighbour.node].push_back(current.node);
	}
}

int main(){
	int nodes, edges, u, v, w;
	fin >> nodes >> edges;
	for(int i = 0; i < edges; ++i){
		fin >> u >> v >> w;
		--u, --v;

		edge[u].push_back({v, w});
		edge[v].push_back({u, w});
	}

	Prim(nodes);

	fout << apm_cost << '\n' << nodes << '\n';
	for(int node = 0; node < nodes; ++node)
		for(auto neighbour: children[node])
			fout << node + 1 << ' ' << neighbour + 1 << '\n';

	fin.close();
	fout.close();
	return 0;
}