// Prims.java

import java.util.LinkedList;
import java.util.List;

// used for making our graph and also our heap
class Edge {
    public int source;
    public int dest;
    public int weight;

    public Edge(int source, int dest, int weight) {
        this.source = source;
        this.dest = dest;
        this.weight = weight;
    }
    
    // this is so we can search for edges by destination
    public boolean equals(Object other) {
        Edge otherEdge = (Edge) other;
        return source == otherEdge.source && dest == otherEdge.dest;
    }
}

// an incidence list based graph
class Graph<Type> { 
    // an array of linked lists, each list holds the edges out of that node
    private LinkedList<Edge>[] nodes;
    
    // this stores the values being stored by this graph
    private Type[] values;

    // the size of the graph
    private int size;

    // construct an empty graph of a given size
    public Graph(int size) {
        this.size = size;

        // becasue this is a generic array, we need to cast
        nodes = (LinkedList<Edge>[]) new LinkedList[size];
        for (int i = 0; i < size; i++) {
            nodes[i] = new LinkedList<Edge>();
        }
        
        // make space for the values (and ignore the cast warning)
        @SuppressWarnings("unchecked")
        Type[] values = (Type[]) new Object[size];
        this.values = values;
    }

    // returns the list of edges associated with a node
    public List<Edge> getNeighbors(int node) {
        return nodes[node];
    }

    // insert an edge
    public void insertEdge(int from, int to, int cost) {
        nodes[from].add(new Edge(from, to, cost));
        nodes[to].add(new Edge(to, from, cost));
    }

    // insert an edge by value
    public void insertEdge(Type from, Type to, int cost) {
        int fromIndex = lookup(from);
        int toIndex = lookup(to);
        insertEdge(fromIndex, toIndex, cost);
    }
    
    // return the cost of an edge or 0 for none
    public int getCost(int from, int to) {
        // search the linked list for this edge
        int index = nodes[from].indexOf(new Edge(from, to, -1));
        
        // if -1, return MAX
        // else return the cost of that edge
        if (index == -1) {
            return Integer.MAX_VALUE;
        } else {
            return nodes[from].get(index).weight;
        }
    }

    // lookup a node number by value
    public int lookup(Type value) {
        for (int i = 0; i < size; i++) {
            if (values[i] != null && values[i].equals(value)) {
                return i;
            }
        }
        return -1;
    }


    // add a node's data to the graph
    public void setValue(int node, Type value) {
        values[node] = value;
    }

    // return the size of the graph
    public int getSize() {
        return size;
    }

    // get the value of a node
    public Type getValue(int index) {
        return values[index];
    }
}

// this heap is used to keep track of the nodes used by
// Prim's algorithm -- it stores the node index and the
// edge cost to add it into the MST
class EdgeHeap {
    private Edge[] array;
    private int count;

    // start with nothing
    public EdgeHeap(int size) {
        array = new Edge[size];
        count = 0;
    }

    // find the left child of a node
    public int left(int node) {
        return 2 * node  + 1;
    }

    // find the right child of a node
    public int right(int node) {
        return 2 * node  + 2; 
    }

    // find the parent of a node
    public int parent(int node) {
        return (node - 1) / 2;
    }

    // insert an item into the heap
    public void insert(int source, int dest, int weight) {
        // put item in next available slot
        array[count] = new Edge(source, dest, weight);

        // keep indices into node we inserted and its parent
        int node = count;
        int p = parent(node);

        // increment total number of items
        count++;

        // upheap until it's root, or parent is smaller (using cost)
        while ((node != 0) && (array[node].weight < array[p].weight)) {
            // swap values
            Edge temp = array[node];
            array[node] = array[p];
            array[p] = temp;

            // update indices
            node = p;
            p = parent(node);
        }
    }

    // remove the top item -- returns the edge
    public Edge dequeue() {
        // set aside root
        Edge retValue = array[0];

        // move last value to the root
        array[0] = array[count - 1];
        count--;

        // find the children
        int node = 0;
        int l = left(node);
        int r = right(node);

        // while one child is smaller than us
        while ((l < count) && (array[l].weight < array[node].weight) ||
               (r < count) && (array[r].weight < array[node].weight)) {

            // find smallest child
            int min;
            if (l < count && r < count) {
                if (array[l].weight < array[r].weight) {
                    min = l;
                } else {
                    min = r;
                }
            } else {
                min = l;
            }

            // swap with the smallest child
            Edge temp = array[node];
            array[node] = array[min];
            array[min] = temp;

            // update indices
            node = min;
            l = left(node);
            r = right(node);
        }
        
        // return orig. root
        return retValue;
    }

    // return the smallest item in the heap
    public Edge min() {
        return array[0];
    }

    public int getCount() {
        return count;
    }
}

public class Prims {
    public static Graph<String> findMST(Graph<String> graph) {
        // the size of the graph (and eventual MST)
        int N = graph.getSize();

        // create a graph for the MST with the same number of nodes
        Graph<String> mst = new Graph<String>(N);

        // create the heap of nodes 
        EdgeHeap heap = new EdgeHeap(N*N);

        // keep track of which nodes we have put into MST
        boolean[] inMST = new boolean[N];
        for (int i = 0; i < N; i++) {
            inMST[i] = false;
        }

        // choose node 0 to be in the MST first and add it to graph
        int size = 1;
        mst.setValue(0, graph.getValue(0));
        inMST[0] = true;

        // add the edges of node 0 into the heap
        List<Edge> neighbors = graph.getNeighbors(0);
        for (Edge e : neighbors) {
            heap.insert(e.source, e.dest, e.weight);
        }

        // while the MST doesn't have all the nodes yet
        while (size < N) {
            // take out the edge with next lowest cost
            Edge next = heap.dequeue();
            int dest = next.dest;

            // if it is already in the MST, move on
            if (inMST[dest]) continue;

            // put destination into the MST
            mst.setValue(dest, graph.getValue(dest));
            inMST[dest] = true;
            size++;

            // add this edge to the tree
            mst.insertEdge(next.source, dest, graph.getCost(next.source, dest));
            System.out.println("Adding edge between " + graph.getValue(next.source)
                    + " and " + graph.getValue(dest));

            // go through all neighbors of this node
            neighbors = graph.getNeighbors(dest);
            for (Edge e : neighbors) {
                if (!inMST[e.dest]) {
                    heap.insert(e.source, e.dest, e.weight);
                }
            }
        }
    
        // return the computer MST
        return mst;
    }

    public static void main(String args[]) {
        // create a graph with the cities
        Graph<String> graph = new Graph<String>(11);
        graph.setValue(0, "Alexandria");
        graph.setValue(1, "Blacksburg");
        graph.setValue(2, "Charlottesville");
        graph.setValue(3, "Danville");
        graph.setValue(4, "Fredericksburg");
        graph.setValue(5, "Harrisonburg");
        graph.setValue(6, "Lynchburg");
        graph.setValue(7, "Newport News");
        graph.setValue(8, "Richmond");
        graph.setValue(9, "Roanoke");
        graph.setValue(10, "Virginia Beach");

        graph.insertEdge("Alexandria", "Fredericksburg", 50);
        graph.insertEdge("Fredericksburg", "Richmond", 60);
        graph.insertEdge("Richmond", "Charlottesville", 70);
        graph.insertEdge("Richmond", "Lynchburg", 110);
        graph.insertEdge("Richmond", "Danville", 145);
        graph.insertEdge("Richmond", "Newport News", 70);
        graph.insertEdge("Newport News", "Virginia Beach", 35);
        graph.insertEdge("Virginia Beach", "Danville", 210);
        graph.insertEdge("Danville", "Lynchburg", 70);
        graph.insertEdge("Lynchburg", "Charlottesville", 70);
        graph.insertEdge("Lynchburg", "Roanoke", 65);
        graph.insertEdge("Roanoke", "Blacksburg", 40);
        graph.insertEdge("Blacksburg", "Harrisonburg", 140);
        graph.insertEdge("Harrisonburg", "Alexandria", 135);
        graph.insertEdge("Harrisonburg", "Charlottesville", 50);

        // compute minimal spanning tree
        findMST(graph);  
    }
}