import java.util.*;

// shortest path
public class GraphAdjacentMatrix_MST_kruskal {
    public static int empty = -1;

    private int adjMatrix[][];
    private int numVertices;
    private boolean addedForMst[];

    public GraphAdjacentMatrix_MST_kruskal(int numVertices) {
        this.numVertices = numVertices;

        adjMatrix = new int[numVertices][numVertices];
        for (int i = 0; i < adjMatrix.length; i++) {
            for (int j = 0; j < adjMatrix[i].length; j++) {
                adjMatrix[i][j] = empty;
            }
        }

        addedForMst = new boolean[numVertices];
    }

    public int getNumVertices(){
        return numVertices;
    }

    public int getWeight(int i, int j){
        return adjMatrix[i][j];
    }

    public void setEdge(int i, int j, int weight) {
        adjMatrix[i][j] = weight;
//        adjMatrix[j][i] = true; // directed

        addedForMst[i] = true;
        addedForMst[j] = true;
    }

    public void delEdge(int i, int j) {
        adjMatrix[i][j] = empty;
//        adjMatrix[j][i] = false; // directed
    }

    public boolean isEdge(int i, int j) {
        return (adjMatrix[i][j] != -1);
    }

    public int first(int v1){
        int[] ary = this.adjMatrix[v1];
        int firstNeighborIndex = this.numVertices; // return n if none
        for (int i = 0; i < ary.length; i++) {
            if (ary[i] != empty){
                firstNeighborIndex = i;
                break;
            }
        }
        return firstNeighborIndex;
    }

    public int next(int v1, int v2){
        int[] ary = this.adjMatrix[v1];
        int nextNeighborIndex = this.numVertices; // return n if none
        for (int i = v2 + 1; i < ary.length; i++) {
            if (ary[i] != empty){
                nextNeighborIndex = i;
                break;
            }
        }
        return nextNeighborIndex; // return n if none
    }

    public class EdgeElemnt implements Comparable{
        public Integer v1;
        public Integer v2;
        public Integer edgeWeight;

        public EdgeElemnt(Integer v1, Integer v2, Integer edgeWeight) {
            this.v1 = v1;
            this.v2 = v2;
            this.edgeWeight = edgeWeight;
        }

        @Override
        public boolean equals(Object obj)
        {
            if(this == obj)
                return true;
            if(obj == null || obj.getClass()!= this.getClass())
                return false;

            // type casting of the argument.
            EdgeElemnt geek = (EdgeElemnt) obj;
            return  (geek.v1 == this.v1 && geek.v2 == this.v2) || (geek.v1 == this.v2 && geek.v2 == this.v1);
        }

        @Override
        public int hashCode() {
            return Objects.hash(v1 + v2);
        }

        @Override
        public int compareTo(Object o) {
            EdgeElemnt o1 = (EdgeElemnt) o;
            return this.edgeWeight.compareTo(o1.edgeWeight);
        }

        @Override
        public String toString() {
            return " " + edgeWeight;
//            return "EdgeElemnt{" +
//                    "fromV=" + fromV +
//                    ", toV='" + toV +
//                    ", edgeWeight=" + edgeWeight +
//                    '}';
        }
    }



    private boolean isReachable(int vDest){

        for (int i = 0; i < numVertices; i++) {
            if (this.isEdge(i, vDest)) return true;
        }

        return false;

    }

    private Set<EdgeElemnt> extractEdgeElement(int vStart){
        boolean[] visitedTmp = new boolean[numVertices];
        Queue<Integer> queue = new LinkedList<>();
        Set<EdgeElemnt> set = new HashSet<>();

        queue.add(vStart);

        while(queue.peek() != null){
            Integer currV = queue.remove();

            if (visitedTmp[currV]) continue;
            visitedTmp[currV] = true;

            for (int nextV = this.first(currV); nextV < this.numVertices; nextV = this.next(currV, nextV)) {
                set.add(new EdgeElemnt(currV, nextV, this.getWeight(currV, nextV))); // go check how equals() and hash() are defined for the EdgeElemnt class
                queue.add(nextV);
            }
        }

        return set;
    }

    private boolean contains(Integer v) {
        return addedForMst[v];
    }

    public static class EdgeWeightComparrator implements Comparator<EdgeElemnt> {
        @Override
        public int compare(EdgeElemnt ee1, EdgeElemnt ee2) {
            return ee1.edgeWeight.compareTo(ee2.edgeWeight);
        }
    }

    static <E extends Comparable<? super E>> void heapsort(E[] ary) { // Heapsort
        MaxHeap<E> H = new MaxHeap<>(ary, ary.length, ary.length);
    }

    public String toString() {
        StringBuilder s = new StringBuilder();
        for (int i = 0; i < numVertices; i++) {
            s.append(i + ": ");
            for (int j : adjMatrix[i]) {
                s.append((j != empty?1:0) + " ");
            }
            s.append("\n");
        }
        return s.toString();
    }

    //////////// MST: kruskal algo ////////////

    public static GraphAdjacentMatrix_MST_kruskal kruskal(GraphAdjacentMatrix_MST_kruskal gam) {

        Set<EdgeElemnt> edgeElemnts = gam.extractEdgeElement(1);

        GraphAdjacentMatrix_MST_kruskal MST = kruskal(edgeElemnts, gam.getNumVertices());
        return MST;
    }

    // calculate the shortest distances between the startingVertice to all other vertices
    public static GraphAdjacentMatrix_MST_kruskal kruskal(Set<EdgeElemnt> edgeElemnts, int numOfVertices) {

        // result
        GraphAdjacentMatrix_MST_kruskal MST = new GraphAdjacentMatrix_MST_kruskal(numOfVertices);
        int totalWeight = 0;

        // use minHeap to kick out the smallest edge each round
        PriorityQueue<EdgeElemnt> minHeap = new PriorityQueue<>(new EdgeWeightComparrator());
        minHeap.addAll(edgeElemnts);

        // Remove items from the Priority Queue (DEQUEUE)
        while (!minHeap.isEmpty()) {
            EdgeElemnt ee = minHeap.remove();

            // check if vertices are in the MST already
            if (MST.contains(ee.v1) && MST.contains(ee.v2)) {
                if (MST.isEdge(ee.v1, ee.v2)) continue; // if the direct edge is already added in MST, skip it.
                // check if they are connected indirectly
                if (MST.isReachable(ee.v2)) continue; // no need to add as long as they are connected indirectly
            }

            // add vertices to the MST
            MST.setEdge(ee.v1, ee.v2, ee.edgeWeight);
            MST.setEdge(ee.v2, ee.v1, ee.edgeWeight);

            totalWeight += ee.edgeWeight;

            System.out.printf("MST(total_w: %d):%n%s %n", totalWeight, MST);
        }

        System.out.printf("MST(total_w: %d):%n%s %n", totalWeight, MST);
        return MST;
    }


    ////////// UTIL ///////////

    private static void printArray(int[] ary){
        for (int i = 0; i < ary.length; i++) {
            System.out.printf("%4d" , ary[i]);
        }
        System.out.println();
    }



    public static void main(String[] args) {
        GraphAdjacentMatrix_MST_kruskal gam = new GraphAdjacentMatrix_MST_kruskal(7);
        // undirected graph
        gam.setEdge(1,2, 5); gam.setEdge(2,1, 5);
        gam.setEdge(1,3, 6); gam.setEdge(3,1, 6);
        gam.setEdge(1,4, 4); gam.setEdge(4,1, 4);
        gam.setEdge(2,3, 1); gam.setEdge(3,2, 1);
        gam.setEdge(2,4, 2); gam.setEdge(3,4, 2);
        gam.setEdge(3,4, 2); gam.setEdge(4,3, 2);
        gam.setEdge(3,5, 5); gam.setEdge(5,3, 5);
        gam.setEdge(3,6, 3); gam.setEdge(6,3, 6);
        gam.setEdge(4,6, 4); gam.setEdge(6,4, 4);
        gam.setEdge(5,6, 4); gam.setEdge(6,5, 4);


        gam.kruskal(gam);

        System.out.println();
    }

}