Implementation of Kurskal's Algorithm using Union-Find

Kruskal.java

import java.util.Arrays;
import java.util.Comparator;

public class Kruskal {

  /**
  * Driver method
  */
  public static void main(String[] args) {
    Graph graph = new Graph(7, 13);
    graph.addEdge(0, 1, 6);
    graph.addEdge(0, 2, 6); 
    graph.addEdge(0, 3, 5); 
    graph.addEdge(1, 3, 8); 
    graph.addEdge(1, 4, 5); 
    graph.addEdge(2, 3, 8); 
    graph.addEdge(2, 5, 3); 
    graph.addEdge(3, 4, 9); 
    graph.addEdge(3, 5, 4); 
    graph.addEdge(5, 6, 1); 
    graph.addEdge(4, 6, 1); 
    graph.addEdge(3, 6, 2); 
    graph.addEdge(2, 6, 1);

    Kruskal k = new Kruskal();
    int[][] result = k.MST(graph);
    for(int[] res : result) {
      System.out.println(res[0] + " " + res[1] + " " + res[2]);
    } 
  } // main

  //-----------------------------------------------------------------------------------------------
  
  /**
   * Computes the MST of given connected graph.
   * Returns array of N-1 edges that belongs to the MST.
   * Assumes that the input graph is connected.
   *
   * @param Graph graph whose MST is to be computed.
   * @return Array of edges belonging to the MST
   *
   * Algorithm:
   *  1. Create a sorted queue of Edges sorted by their weight
   *  2. Greedily Pick lightest edge and include it in the result if it does not create cycle
   *  3. When there are N-1 edges return the edges
   *
   */
  public int[][] MST(Graph g) {
    // Initialize container to hold result edges. There will be N-1 edges
    int[][] result = new int[g.N - 1][3];
    int cursor = 0; // cursor to point to the index to insert an Edge
    
    // Initialize Union-Find data structure, that will assist to know
    // if two vertex are already connected or not
    WQuickUnion uf = new WQuickUnion(g.N);

    // Sort Graph Edges in non-decreasing order of the edge weight
    g.sortEdges();
    int[] edge = g.getNextEdge();
    while(edge != null) {
      if (!uf.find(edge[0], edge[1])) { // check if source and target nodes are connected or not
        result[cursor++] = edge;
        uf.union(edge[0], edge[1]);
        if (cursor == g.N) break; // we already have N-1 Tree edges   
      }
      edge = g.getNextEdge();
    }
    return result;
  } // MST

} // Kruskal


//=================================================================================================

/**
 * Class to represent a graph
 */
class Graph {
  int N; // Number of vertices. Vertices will be numbered 0, 1, 2, 3, ....
  int[][] E; // Edges in Graph. Each edge will be 3 element array representing source vertex, target vertex and the weight.
  int cursor = 0; // cursor for inserting and retrieving edges

  //-----------------------------------------------------------------------------------------------

  /**
    * Constructor
    * @param N number of vertices in the graph
    * @param E number of edges in the graph
    *
    */ 
  public Graph(int N, int E) {
    this.N = N;
    this.E = new int[E][3];
  } // Graph


  //-----------------------------------------------------------------------------------------------

  /**
    * Adds an edge to the graph
    * 
    * @param s source node
    * @param t target node
    * @param w weight of edge
    *
    */ 
  public void addEdge(int s, int t, int w) {
    this.E[cursor++] = new int[]{s, t, w};
  } // addEdge

  //-----------------------------------------------------------------------------------------------
  
  /**
    * Sort the edges of the graph in increasing order of weights.
    *
    */ 
  public void sortEdges() {
    Arrays.sort(this.E, new Comparator<int[]>() {
      @Override
      public int compare(int[] o1, int[] o2) {
        return o1[2] - o2[2];
      } // compare
    });
    this.cursor = 0;
  } // sortEdges

  //-----------------------------------------------------------------------------------------------
  
  /**
    * Get next edge as pointed by cursor. 
    * Cursor will be reset if sortEdges() is called.
    * @return returns current edge pointed by cursor
    */ 
  public int[] getNextEdge() {
    if(cursor == E.length) return null;
    return E[cursor++];
  } // getNextEdge

} // Graph


//=================================================================================================

/**
 * Class to represent Union-Find Data structure
 * This is Weighted Quick Union version of the data structure.
 *
 */
class WQuickUnion {
  private int[] id; // id[i] represent parent of the element at index i
  private int[] size; //size[i] represent number of nodes element at index has

  //-----------------------------------------------------------------------------------------------
  
  /**
   * Constructor to initilize the data structure.
   * @param n number of elements.
   *
   */ 
  public WQuickUnion(int n) {
    id = new int[n];
    size = new int[n];
    for (int i = 0; i < n; i++) {
      id[i] = i;
      size[i] = 1;
    }
  } // WQuickUnion

  //-----------------------------------------------------------------------------------------------
  
  /**
   * Identifies whether two nodes p and q are connected or not
   * @param p first node
   * @param q second node
   * @return True if the two nodes are connected otherwise False
   */ 
  public boolean find(int p, int q) {
    return root(p) == root(q);
  } // find

  //-----------------------------------------------------------------------------------------------
  
  /**
   * Connects two nodes in the network
   * @param p first node
   * @param q second node
   */ 
  public void union(int p, int q) {
    int i = root(p);
    int j = root(q);
    if (size[i] < size[j]) {
      id[i] = j;
      size[j] += size[i];
    } else {
      id[j] = i;
      size[i] += size[j];
    }
    
  } // union

  //-----------------------------------------------------------------------------------------------

  /**
   * Finds the root of current element
   * @param p current element index
   * @return index of the root of input index
   */ 
  private int root(int p) {
    while(p != id[p]) {
      id[p] = id[id[p]]; // Path Compression. Keeps tree almost completely flat. Will still work if this line is removed.
      p = id[p];
    }
    return p;
  } // root

} // WQuickUnion