Cod sursa(job #2711914)

Utilizator vnedelcuVictor Andrei Nedelcu vnedelcu Data 24 februarie 2021 21:10:43
Problema Arbore partial de cost minim Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 5.1 kb
#include <stdio.h>
#include <vector>
#include <limits.h>

using namespace std;

const int INF = INT_MAX;

class MinHeap {
private:
    int heap_size;
    vector<pair<int, pair<int, int> > > cost_edge_pairs;
    vector<int> heap_positions;

    int father(int heap_pos) {return heap_pos >> 1;}
    int left_son(int heap_pos) {return heap_pos << 1;}
    int right_son(int heap_pos) {return (heap_pos << 1) | 1;}
    void up_heap(int heap_pos);
    void down_heap(int heap_pos);

public:
    MinHeap(int capacity);
    void insert(pair<int, int> edge, int value);
    void remove(int heap_pos);
    void decrease_cost(int heap_pos, int new_cost, pair<int, int> new_edge);
    int get_heap_pos(int node) {return heap_positions[node];}
    pair<int, pair<int, int> > top() {return cost_edge_pairs[1];}
};

MinHeap::MinHeap(int capacity) {
    cost_edge_pairs.resize(capacity + 1);
    heap_positions.resize(capacity);
    heap_size = 0;
}

void MinHeap::insert(pair<int, int> edge, int value) {
    ++heap_size;
    cost_edge_pairs[heap_size] = {value, edge};
    heap_positions[edge.second] = heap_size;
    up_heap(heap_size);
}

void MinHeap::remove(int heap_pos) {
    cost_edge_pairs[heap_pos] = cost_edge_pairs[heap_size];
    heap_positions[cost_edge_pairs[heap_size].second.second] = heap_pos;
    --heap_size;
    if (heap_pos > 1 && cost_edge_pairs[heap_pos].first < cost_edge_pairs[father(heap_pos)].first) {
        up_heap(heap_pos);
    } else {
        down_heap(heap_pos);
    }
}

void MinHeap::decrease_cost(int heap_pos, int new_cost, pair<int, int> new_edge) {
    cost_edge_pairs[heap_pos].first = new_cost;
    cost_edge_pairs[heap_pos].second = new_edge;
    up_heap(heap_pos);
}

void MinHeap::up_heap(int heap_pos) {
    pair<int, pair<int, int> > temp_pair = cost_edge_pairs[heap_pos];

    int f = father(heap_pos);
    while (f > 0 && cost_edge_pairs[f].first > temp_pair.first) {
        cost_edge_pairs[heap_pos] = cost_edge_pairs[f];
        heap_positions[cost_edge_pairs[f].second.second] = heap_pos;
        heap_pos = f;
        f = father(heap_pos);
    }

    cost_edge_pairs[heap_pos] = temp_pair;
    heap_positions[temp_pair.second.second] = heap_pos;
}

void MinHeap::down_heap(int heap_pos) {
    while (true) {
        int min_son = 0;
        if (left_son(heap_pos) <= heap_size) {
            min_son = left_son(heap_pos);
            if (
                right_son(heap_pos) <= heap_size &&
                cost_edge_pairs[min_son].first > cost_edge_pairs[right_son(heap_pos)].first
            ) {
                min_son = right_son(heap_pos);
            }
        }

        if (min_son && cost_edge_pairs[min_son].first < cost_edge_pairs[heap_pos].first) {
            swap(cost_edge_pairs[min_son], cost_edge_pairs[heap_pos]);
            swap(
                heap_positions[cost_edge_pairs[min_son].second.second],
                heap_positions[cost_edge_pairs[heap_pos].second.second]
            );
            heap_pos = min_son;
        } else {
            break;
        }
    }
}

int main() {
    FILE * f;
    int n, m, x, y, c;

    // reading data and building the graph
    f = fopen("apm.in", "r");
    fscanf(f, "%d%d", &n, &m);
    vector<vector<pair<int, int> > > graph(n);
    for (int i = 0; i < m; ++i) {
        fscanf(f, "%d%d%d", &x, &y, &c);
        graph[x - 1].push_back({y - 1, c});
        graph[y - 1].push_back({x - 1, c});
    }
    fclose(f);

    // prim's algorithm
    vector<pair<int, int> > apm(n - 1);
    vector<bool> visited(n, false);
    vector<int> costs(n, INF);
    MinHeap minheap(n);

    visited[0] = true;
    for (int i = 0; i < graph[0].size(); ++i) {
        minheap.insert({0, graph[0][i].first}, graph[0][i].second);
        costs[graph[0][i].first] = graph[0][i].second;
    }

    int total_cost = 0;
    for (int i = 0; i < n - 1; ++i) {
        pair<int, pair<int, int> > top_heap = minheap.top();
        minheap.remove(1);

        int new_apm_node = top_heap.second.second;
        apm[i] = top_heap.second;
        total_cost += top_heap.first;
        visited[new_apm_node] = true;

        for (int j = 0; j < graph[new_apm_node].size(); ++j) {
            int neighbor = graph[new_apm_node][j].first;
            int cost = graph[new_apm_node][j].second;

            if (!visited[neighbor]) {
                if (costs[neighbor] == INF) {
                    minheap.insert({new_apm_node, neighbor}, cost);
                    costs[neighbor] = cost;
                } else if (cost < costs[neighbor]) {
                    minheap.decrease_cost(
                        minheap.get_heap_pos(neighbor),
                        cost,
                        {new_apm_node, neighbor}
                    );
                    costs[neighbor] = cost;
                }
            }
        }
    }

    // printing the output
    f = fopen("apm.out", "w");
    fprintf(f, "%d\n%d\n", total_cost, n - 1);
    for (int i = 0; i < n - 1; ++i) {
        fprintf(f, "%d %d\n", apm[i].first + 1, apm[i].second + 1);
    }
    fclose(f);

    return 0;
}