Cod sursa(job #1741794)

Utilizator AlexandruValeanuAlexandru Valeanu AlexandruValeanu Data 15 august 2016 09:51:25
Problema Arbore partial de cost minim Scor 0
Compilator java Status done
Runda Arhiva educationala Marime 5.47 kb
import java.io.*;
import java.util.*;

public class Main {

    public static void main(String[] args) throws IOException {
        InputReader in = new InputReader(new FileInputStream("apm.in"));

        int N = in.nextInt();
        int M = in.nextInt();

        WeightedGraph weightedGraph = new WeightedGraph(N);

        for (int i = 0; i < M; i++) {
            int x = in.nextInt();
            int y = in.nextInt();
            int w = in.nextInt();
            weightedGraph.addEdge(x, y, w);
            weightedGraph.addEdge(y, x, w);
        }

        PrimMST mst = new PrimMST(weightedGraph);
        List<Edge> treeEdges = mst.minimumSpanningTree();

        int totalCost = 0;

        for (Edge edge: treeEdges){
            totalCost += edge.cost;
        }

        PrintWriter out = new PrintWriter(new FileOutputStream("apm.out"));

        out.println(totalCost + "\n" + (N - 1));

        for (Edge edge: treeEdges) {
            out.println(edge.x + " " + edge.y);
        }

        out.close();
    }

    static class Node implements Comparable<Node>{
        int distance;
        int node;

        public Node(int distance, int node) {
            this.distance = distance;
            this.node = node;
        }

        @Override
        public int compareTo(Node o) {
            return Integer.compare(distance, o.distance);
        }
    }

    static class PrimMST {
        private WeightedGraph weightedGraph;
        private int N;
        private int[] distance, father;
        private boolean[] visited;

        PrimMST(WeightedGraph WG){
            this.weightedGraph = WG;
            this.N = WG.getN();
            this.distance = new int[N + 1];
            this.father = new int[N + 1];
            this.visited = new boolean[N + 1];
        }

        ArrayList<Edge> minimumSpanningTree(){
            ArrayList<Edge> treeEdges = new ArrayList<>();
            Queue<Node> minHeap = new PriorityQueue<>();

            Arrays.fill(distance, Integer.MAX_VALUE);
            Arrays.fill(father, -1);
            Arrays.fill(visited, false);

            distance[1] = 0;
            minHeap.add(new Node(distance[1], 1));

            while (!minHeap.isEmpty()){
                Node topNode = minHeap.remove();
                int node = topNode.node;

                if (distance[node] != topNode.distance)
                    continue;

                visited[node] = true;

                if (father[node] != -1)
                    treeEdges.add(new Edge(node, father[node], distance[node]));

                for (int p = weightedGraph.getHead(node); p != WeightedGraph.NIL; p = weightedGraph.getNext(p)) {
                    int son = weightedGraph.getNeighbour(p);
                    int cost = weightedGraph.getCost(p);

                    if (!visited[son] && distance[son] > cost){
                        distance[son] = cost;
                        father[son] = node;
                        minHeap.add(new Node(distance[son], son));
                    }
                }
            }

            return treeEdges;
        }
    }

    static class WeightedGraph{
        private static class Edge{
            int node;
            int cost;
            int next;

            Edge(int node, int cost, int next){
                this.node = node;
                this.cost = cost;
                this.next = next;
            }
        }

        final static int NIL = -1;
        private int[] head;
        private ArrayList<Edge> graph;

        private int N;
        private int counter;

        WeightedGraph(int N){
            initialize(N);
        }

        public int getN() {
            return N;
        }

        private void initialize(final int N){
            head = new int[N];
            graph = new ArrayList<>();

            this.N = N;
            this.counter = 0;

            Arrays.fill(head, NIL);
        }

        void addEdge(int x, int y, int w){
            if (!(1 <= x && x <= N)) throw new AssertionError();
            x--; y--;
            graph.add(new Edge(y, w, head[x]));
            head[x] = counter++;
        }

        int getHead(int node){
            if (!(1 <= node && node <= N)) throw new AssertionError();
            node--;
            return head[node];
        }

        int getNext(int p){
            if (!(0 <= p && p < counter)) throw new AssertionError();
            return graph.get(p).next;
        }

        int getNeighbour(int p){
            if (!(0 <= p && p < counter)) throw new AssertionError();
            return graph.get(p).node + 1;
        }

        int getCost(int p){
            if (!(0 <= p && p < counter)) throw new AssertionError();
            return graph.get(p).cost;
        }
    }

    static class InputReader{
        private BufferedReader reader;
        private StringTokenizer tokenizer;

        InputReader(InputStream stream) {
            reader = new BufferedReader(new InputStreamReader(stream), 65536);
            tokenizer = null;
        }

        private String nextToken() {
            while (tokenizer == null || !tokenizer.hasMoreTokens()){
                try {
                    tokenizer = new StringTokenizer(reader.readLine());
                }
                catch (IOException e){
                    throw new RuntimeException(e);
                }
            }

            return  tokenizer.nextToken();
        }

        int nextInt() {
            return Integer.parseInt(nextToken());
        }

        String nextString(){
            return nextToken();
        }
    }
}