Example: Connecting cities - rFronteddu/general_wiki GitHub Wiki

There are N cities numbered from 1 to N.

You are given connections, where each connections[i] = [city1, city2, cost] represents the cost to connect city1 and city2 together. (A connection is bidirectional: connecting city1 and city2 is the same as connecting city2 and city1.)

Return the minimum cost so that for every pair of cities, there exists a path of connections (possibly of length 1) that connects those two cities together. The cost is the sum of the connection costs used. If the task is impossible, return -1.

Solution

We can use Kruskal's Algorithm:

1. Sort edges in non-decreasing order.
2. Implement UnionFind structure to manage connected component (to determine whether adding an edge will form a cycle)
3. Process edges, add an edge if it connects two previously unconnected components. If it can, add its cost to the total cost and union the two components.
4. Check Connectivity, number of edges in the MST must be N-1 otherwise it is not fully connected.

public static int minimumCost(int N, int[][] connections) {
    // we can use Kruskal to compute the MST
    Arrays.sort(connections, (a,b) -> Integer.compare(a[2], b[2]));
    
    UnionFind unionFind = new UnionFind(N);

    int totalCost = 0;
    int edgeUsed = 0;

    // process each edge
    for (int[] connection : connections) {
        // the -1 is to convert in 0-based index
        int city1 = connection[0] - 1;
        int city2 = connection[1] - 1;
        int cost = connection[2];


        // check if city1 and 2 are in different components
        if (unionFind.find(city1) != unionFind.find(city2)) {
            // add edge to MST
            unionFind.union(city1, city2);
            totalCost += cost;
            edgeUsed++;
        }
    }

    // check if all cities are connected
    if (edgeUsed == N - 1) {
        return totalCost;
    }
    return -1;
}

static class UnionFind {
    private int[] parent;
    private int[] rank;

    public UnionFind(int size) {
        parent = new int[size];
        rank = new int[size];
        for (int i = 0; i < size; i++) {
            // each node is its own parent
            parent[i] = i;
            // initial rank
            rank[i] = 1;
        }
    }

    public int find(int city) {
        if (parent[city] != city) {
            // path compression
            parent[city] = find(parent[city]);
        }
        return parent[city];
    }

    public void union(int city1, int city2) {
        int root1 = find(city1);
        int root2 = find(city2);
        
        if (root1 != root2) {
            // union by rank
            if (Rank[root1] > rank[root2]) {
                parent[root2] = root1;
            } else if (rank[root1] < rank[root2]) {       
                parent[root1] = root2;
            } else {
                parent[root2] = root1;
                rank[root1]++;
            }
        }
    }
}