1200

src/p2.cpp

@@ -1,16 +1,10 @@
 #include <iostream>
 #include <vector>
 
-#define FINISHED 0
-#define VISITING 1
-#define NOTVISITED 2
-
-#define LOG(x) std::cout << x << std::endl
-
 struct tree_node_t {
     tree_node_t* p;
     size_t rank;
-    int id;
+    size_t id;
 };
 
 struct edge_t {
@@ -21,34 +15,34 @@ struct edge_t {
 
 struct heap_t {
     int size;
-    std::vector<edge_t> edges;
+    std::vector<edge_t*> edges;
 };
 
 static inline int parent(int i) {
-    return i >> 1;
+    return (i >> 1);
 }
 
 static inline int left(int i) {
-    return i << 1;
+    return (i << 1);
 }
 
 static inline int right(int i) {
     return (i << 1) + 1;
 }
 
 void swap(heap_t* heap, int i, int j) {
-    edge_t aux = heap->edges[i];
+    edge_t* aux = heap->edges[i];
     heap->edges[i] = heap->edges[j];
     heap->edges[j] = aux;
 }
 
 void max_heapify(heap_t* heap, int i) {
     int l = left(i), r = right(i), largest = i;
 
-    if (l < heap->size && heap->edges[l].weight > heap->edges[i].weight)
+    if (l < heap->size + 1 && heap->edges[l]->weight > heap->edges[i]->weight)
         largest = l;
 
-    if (r < heap->size && heap->edges[r].weight > heap->edges[largest].weight)
+    if (r < heap->size + 1 && heap->edges[r]->weight > heap->edges[largest]->weight)
         largest = r;
 
     if (largest != i) {
@@ -57,9 +51,37 @@ void max_heapify(heap_t* heap, int i) {
     }
 }
 
+void heap_increase_key(heap_t *heap, int i, size_t key) {
+    if (heap->edges[i]->weight == key)
+        return;
+
+    heap->edges[i]->weight = key;
+
+    while (i > 1 && heap->edges[parent(i)]->weight < heap->edges[i]->weight) {
+        swap(heap, i, parent(i));
+        i = parent(i);
+    }
+}
+
+void max_heap_insert(heap_t* heap, edge_t* edge) {
+    heap->size++;
+    size_t key = edge->weight;
+    edge->weight = 0;
+    heap->edges.push_back(edge);
+    heap_increase_key(heap, heap->size, key);
+}
+
+edge_t* extract_max(heap_t* heap) {
+    edge_t* max = heap->edges[1];
+    heap->edges[1] = heap->edges[heap->size];
+    heap->size--;
+    max_heapify(heap, 1);
+    return max;
+}
+
 void build_max_heap(heap_t* heap) {
-    size_t size = heap->size >> 1;
-    for (size_t i = size; i > 0; i--)
+    int size = heap->size >> 1;
+    for (int i = size; i > 0; i--)
         max_heapify(heap, i);
 }
 
@@ -88,56 +110,56 @@ void node_union(tree_node_t* x, tree_node_t* y) {
     link(find_set(x), find_set(y));
 }
 
-unsigned long int get_maximum_cost_spanning_tree(heap_t* heap, std::vector<tree_node_t*>& nodes) {
+unsigned long int get_maximum_cost_spanning_tree(heap_t* heap) {
     unsigned long int result = 0;
-
-    for (auto & node : nodes)
-        make_set(node);
 
-    for (int i = 0; i < heap->size; i++) {
-        if (find_set(heap->edges[i].u) != find_set(heap->edges[i].v)) {
-            node_union(heap->edges[i].u, heap->edges[i].v);
-            result += heap->edges[i].weight;
+    while (heap->size > 0) {   
+        edge_t* edge = extract_max(heap);
+        if (find_set(edge->u) != find_set(edge->v)) {
+            node_union(edge->u, edge->v);
+            result += edge->weight;
         }
+        delete edge;
     }
 
     return result;
 }
 
-unsigned long int compute_graph(heap_t* heap, std::vector<tree_node_t*> &nodes) {
-    unsigned long int result = 0;
-    result += get_maximum_cost_spanning_tree(heap, nodes);
-    return result;
-}
+void read_input (heap_t *heap, std::vector<tree_node_t*> &nodes) {
+    int v_count, e_count;
 
+    scanf("%d", &v_count);
+    scanf("%d", &e_count);
 
-int main(void) {
-    int v_count, e_count;
-    std::cin >> v_count >> e_count;
-
-    heap_t* heap = new heap_t();
-    std::vector<tree_node_t*> nodes = std::vector<tree_node_t*>(v_count, nullptr);;
-    for (int i = 0; i < e_count; i++) {
-        int id1, id2;
-        size_t weight;
-        scanf("%d %d %ld", &id1, &id2, &weight);
-
-        if (nodes[--id1] == nullptr)
-            nodes[id1] = new tree_node_t({.id = id1});
+    heap->edges.push_back(nullptr);
 
-        if (nodes[--id2] == nullptr)
-            nodes[id2] = new tree_node_t({.id = id2});
+    for (int i = 0; i < v_count; i++) {
+        nodes.push_back(new tree_node_t());
+        nodes[i]->id = i;
+        make_set(nodes[i]);
+    }
 
-        heap->edges.push_back(edge_t({.u = nodes[id1], .v = nodes[id2], .weight = weight}));
+    int id1, id2;
+    size_t weight;
+    for (int i = 0; i < e_count; i++) {
+        scanf("%d %d %ld", &id1, &id2, &weight);
+        max_heap_insert(heap, new edge_t({.u = nodes[id1-1], .v = nodes[id2-1], .weight = weight}));
     }
 
     heap->size = e_count;
-    build_max_heap(heap);
+}
 
-    std::cout << compute_graph(heap, nodes) << std::endl;
+int main(void) {
+    heap_t* heap = new heap_t();
+    std::vector<tree_node_t*> nodes;
+    read_input(heap, nodes);
+
+    std::cout << get_maximum_cost_spanning_tree(heap) << std::endl;
 
     for (auto & node : nodes)
         delete node;
 
     delete heap;
+
+    return 0;
 }

tests/public_tests/test4.in

@@ -0,0 +1,6 @@
+5
+4
+1 2 1
+1 4 6
+4 3 7
+3 2 8

tests/public_tests/test4.out

@@ -0,0 +1 @@
+21