(UVA) Expensive Subway - Solution 2

If you want to see another solution for this problem, click here.

I used Prim's algorithm to solve this problem.


import java.io.*;
import java.util.*;

class Main {
    public static HashMap<String, Integer> map;
    public static HashMap<Integer, ArrayList<Edge>> adjList;
    public static int numVisitedStations;
   
    public static Comparator<Edge> costComparator = new Comparator<Edge>() {
        public int compare(Edge e1, Edge e2) {
            return e1.cost-e2.cost;
        }
    };
   
    public static int prim(int start) {
        Queue<Edge> queue = new PriorityQueue<Edge>(20, costComparator);
        Edge e = new Edge(start, 0);
        queue.add(e);
       
        HashSet<Integer> visited = new HashSet<Integer>();
       
        int totalCost = 0;
        while (queue.size() > 0) {
            Edge curr = queue.poll();
            int currStation = curr.station;
            int currCost = curr.cost;
           
            if (visited.contains(currStation)) {
                continue;
            }
            visited.add(currStation);
            numVisitedStations++;
            totalCost += currCost;
           
            ArrayList<Edge> reachStations = adjList.get(currStation);
            for (int i = 0; i < reachStations.size(); i++) {
                queue.add(reachStations.get(i));
            }
        }
       
        return totalCost;
    }
   
    public static void process() throws NumberFormatException, IOException {
        Scanner sc = new Scanner(System.in);
       
        int numStations = sc.nextInt();
        int numConnections = sc.nextInt();
        while (numStations != 0 || numConnections != 0) {
            map = new HashMap<String, Integer>();
            adjList = new HashMap<Integer, ArrayList<Edge>>();
            for (int i = 0; i < numStations; i++) {
                String name = sc.next();
                map.put(name, i);
                adjList.put(i, new ArrayList<Edge>());
            }
           
            for (int i = 0; i < numConnections; i++) {
                String station1 = sc.next();
                String station2 = sc.next();
                int cost = sc.nextInt();
               
                adjList.get(map.get(station1)).add(new Edge(map.get(station2), cost));
                adjList.get(map.get(station2)).add(new Edge(map.get(station1), cost));
            }
           
            int start = map.get(sc.next());

            numVisitedStations = 0;           
            int totalCost = prim(start);

            if (numVisitedStations == numStations) {
                System.out.println(totalCost);
            }
            else {
                System.out.println("Impossible");
            }
           
            numStations = sc.nextInt();
            numConnections = sc.nextInt();
        }
                                       
        return;
    }
   
    public static void main(String[] args) throws NumberFormatException, IOException {
        Main m = new Main();
        m.process();
       
        System.exit(0);
    }
}

class Edge {
    int station;
    int cost;
   
    public Edge(int s, int c) {
        station = s;
        cost = c;
    }
}