Kruskal MST algorithm

Question

I came up with the following code by following Prof. Sedgewick's lecture on Coursera.

Please review this code and let me know if there is anything that I got wrong in implementing Kruskal's algorithm. I'd also like to know what is Big-O complexity of this algorithm.

import java.util.ArrayDeque;
import java.util.PriorityQueue;
import java.util.Queue;

public class KruskalMST {     
private class WeightedEdge {
                    public int from, to, weight;
                    public WeightedEdge(int from, int to, int weight) {
                        this.from = from;
                        this.to = to;
                        this.weight = weight;
                    }

                    @Override
                    public String toString() {
                        StringBuilder sb = new StringBuilder();
                        sb.append("From --> ");
                        sb.append(from+1);
                        sb.append(", to --> ");
                        sb.append(to+1);
                        sb.append(", weight --> ");
                        sb.append(weight);
                        return sb.toString();
                    }
                }

                private class UnionFind {

                    private int capacity;
                    private int[] arr;
                    private int[] size;
                    public UnionFind(int capacity) {
                        this.capacity = capacity;
                        this.arr = new int[capacity];
                        this.size = new int[capacity];

                        for(int i=0; i < capacity; i++) {
                            this.arr[i] = i;
                            this.size[i] = 1;
                        }
                    }


                    private int root(int i) {
                        while(i != arr[i]) {
                            i = arr[arr[i]];
                        }
                        return i;
                    }


                    public void union(int i, int j) {
                        int p = root(i);
                        int q = root(j);

                        if(p != q) {

                            if(this.size[p] <= this.size[q]) {
                                this.size[q] += this.size[p];

                                this.arr[p] = q;
                            }
                            else {

                                this.arr[q] = p;
                                this.size[p] += this.size[q];
                            }

                        }

                    }

                    public boolean connected(int i, int j) {
                         int p = root(i);
                         int q = root(j);
                         return p == q;
                    }

                }

                public Queue<WeightedEdge> findMinCostConnectionToAllCities(int[][] roadNetwork) {


                    int n = roadNetwork.length;
                    PriorityQueue<WeightedEdge> pq = new PriorityQueue<WeightedEdge>(2*n,(WeightedEdge e1, WeightedEdge e2) -> {
                        return e1.weight - e2.weight;
                    });
                    Queue<WeightedEdge> mst = new ArrayDeque<>();
                    UnionFind uf = new UnionFind(n);

                    for(int i=0; i < n; i++) {
                        for(int j=0; j < n; j++) {
                            if(i != j && roadNetwork[i][j] > 0) {

                                WeightedEdge edge = new WeightedEdge(i,j,roadNetwork[i][j]);
                                pq.add(edge);
                            }
                        }
                    }

                    while(!pq.isEmpty() && mst.size() < n-1) {

                        WeightedEdge edge = pq.remove();
                        if(!uf.connected(edge.from, edge.to)) {

                            uf.union(edge.from,edge.to);
                            mst.add(edge);
                        }
                    }

                    return mst;

                }

                public static void main (String[] args) {
                    int[][] city1 = {{0, 1, 2, 3, 4}, 
                                           {1, 0, 5, 0, 7}, 
                                           {2, 5, 0, 6, 0}, 
                                           {3, 0, 6, 0, 0}, 
                                           {4, 7, 0, 0, 0}}; 

                    int[][] city2 = {{0, 1, 1, 100, 0, 0}, 
                            {1, 0, 1, 0, 0, 0}, 
                            {1, 1, 0, 0, 0, 0},    
                            {100, 0, 0, 0, 2, 2}, 
                            {0, 0, 0, 2, 0, 2},   
                            {0, 0, 0, 2, 2, 0}}; 

                    KruskalMST kruskal = new KruskalMST();

                    Queue<WeightedEdge> mst =   kruskal.findMinCostConnectionToAllCities(city2);
                    int totalCost = 0;
                    for(WeightedEdge edge: mst) {
                        totalCost += edge.weight;
                        System.out.println(edge.toString());
                    }

                    System.out.println("Total cost --> " + totalCost);

                }


        }

Answer 1

General

This should probably be three separate classes, unless there's some compelling reason you haven't shared to wrap WeightedEdge and UnionFind inside KruskalMST. If they must be internal, they should be static because they don't rely on context from KruskalMST.

You should use final to indicate classes are not designed for extension and that properties will not change after their initial assignment. This reduces cognitive load on the reader and gives hints to the compiler.

In idiomatic Java, we only declare one variable per line, even if they share a type.

In idiomatic Java, put whitespace between a control flow keyword (for, while and the opening {.

In idomatic Java, else { belongs on the same line as }, not a newline.

It's preferred to include whitespace on both sides of binary operations (+, -, etc).

It's preferred to have whitespace after a ,. Please be consistent.

Don't use abbreviations when naming variables. It makes it harder for readers to understand your code. In general, good names are very helpful for increasing readability.

Methods that return a boolean variable typically begin with a predicate such as is or has. In the case of connected, it's arguable whether isConnected (gramatically incorrect, but standard predicate) is preferable to areConnected (gramatically correct, but unusual predicate). I would argue that both are preferable to connected, which only loosely suggests it returns a boolean value.

WeightedEdge

Member variables should be private unless you have a compelling reason to expose them. This allows your class to control its internal representation of properties without breaking clients if it changes. Use accessor methods to allow clients access the the information. In this case, given that the values are primitives and the class is immutable (its state will not change after object creation), it's not terrible to expose them. However, you definitely should make the values final. Classes should control their own internals.

Since you have a known need for a comparator by weight, it would be cleaner to expose it on WeightedEdge rather than force clients to create it.

UnionFind

capacity is unused outside the constructor and does not need to be stored as an instance variable.

It's typically considered a poor practice to reassign method parameter variables.

We can gain readability at the cost of some performance by letting union() use connected(). This is often an excellent tradeoff, but varies on a case-by-case basis.

A guard clause might make union a bit easier to read.

connected can be simplified to a single line.

Kruskal MST

Your comparator can overflow or underflow for extreme values. Presumably this isn't an issue in your case, but it's something to be aware of. Using Integer.compare(int, int) would be preferable.

You can reduce the number of loop iterations in half by adding both edge directions at the same time.

If you made all these changes, your code might look something like:

import java.util.ArrayDeque;
import java.util.PriorityQueue;
import java.util.Queue;

public final class KruskalMST {

    public Queue<WeightedEdge> findMinimumCostConnectionToAllCities(final int[][] roadNetwork) {

        final int n = roadNetwork.length;
        final PriorityQueue<WeightedEdge> edges =
                new PriorityQueue<WeightedEdge>(2 * n, WeightedEdge.WEIGHT_DESCENDING_ORDER);

        final Queue<WeightedEdge> minimumSpanningTree = new ArrayDeque<>();
        final UnionFind unionFind = new UnionFind(n);

        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                if (roadNetwork[i][j] > 0) {
                    edges.add(new WeightedEdge(i, j, roadNetwork[i][j]));
                }
                if (roadNetwork[j][i] > 0) {
                    edges.add(new WeightedEdge(j, i, roadNetwork[j][i]));
                }
            }
        }

        while (!edges.isEmpty() && minimumSpanningTree.size() < n - 1) {
            final WeightedEdge edge = edges.remove();
            if (!unionFind.areConnected(edge.from, edge.to)) {
                unionFind.union(edge.from, edge.to);
                minimumSpanningTree.add(edge);
            }
        }

        return minimumSpanningTree;

    }

    public static void main(final String[] args) {
        final int[][] city1 = {
                {0, 1, 2, 3, 4},
                {1, 0, 5, 0, 7},
                {2, 5, 0, 6, 0},
                {3, 0, 6, 0, 0},
                {4, 7, 0, 0, 0}};

        final int[][] city2 = {
                {0, 1, 1, 100, 0, 0},
                {1, 0, 1, 0, 0, 0},
                {1, 1, 0, 0, 0, 0},
                {100, 0, 0, 0, 2, 2},
                {0, 0, 0, 2, 0, 2},
                {0, 0, 0, 2, 2, 0}};

        final KruskalMST kruskal = new KruskalMST();
        final Queue<WeightedEdge> minimumSpanningTree = kruskal.findMinimumCostConnectionToAllCities(city2);

        int totalCost = 0;
        for (final WeightedEdge edge: minimumSpanningTree) {
            totalCost += edge.weight;
            System.out.println(edge.toString());
        }
        System.out.println("Total cost --> " + totalCost);
    }

}

---

import java.util.Comparator;

final class WeightedEdge {

    /** Sorts edges from greatest weight to least weight. */
    public static final Comparator<WeightedEdge> WEIGHT_DESCENDING_ORDER = 
            (final WeightedEdge e1, final WeightedEdge e2) -> Integer.compare(e1.weight, e2.weight);

    public final int from;
    public final int to;
    public final int weight;

    public WeightedEdge(final int from, final int to, final int weight) {
        this.from = from;
        this.to = to;
        this.weight = weight;
    }

    @Override
    public String toString() {
        final StringBuilder stringBuilder = new StringBuilder();
        stringBuilder.append("From --> ");
        stringBuilder.append(from + 1);
        stringBuilder.append(", to --> ");
        stringBuilder.append(to + 1);
        stringBuilder.append(", weight --> ");
        stringBuilder.append(weight);
        return stringBuilder.toString();
    }
}

---

final class UnionFind {

    private final int[] arr;
    private final int[] size;

    public UnionFind(final int capacity) {
        this.arr = new int[capacity];
        this.size = new int[capacity];

        for (int i = 0; i < capacity; i++) {
            this.arr[i] = i;
            this.size[i] = 1;
        }
    }

    public void union(final int i, final int j) {
        if (areConnected(i, j)) {
            return;
        }

        final int rootOfI = rootOf(i);
        final int rootOfJ = rootOf(j);

        if (this.size[rootOfI] <= this.size[rootOfJ]) {
            this.size[rootOfJ] += this.size[rootOfI];
            this.arr[rootOfI] = rootOfJ;
        } else {
            this.arr[rootOfJ] = rootOfI;
            this.size[rootOfI] += this.size[rootOfJ];
        }
    }

    public boolean areConnected(final int i, final int j) {
        return rootOf(i) == rootOf(j);
    }


    private int rootOf(final int i) {
        int parent = i;
        while (parent != arr[parent]) {
            parent = arr[arr[parent]];
        }
        return parent;
    }

}