Lintcode 629. Minimum Spanning Tree

Given a list of Connections, which is the Connection class (the city name at both ends of the edge and a cost between them), find some edges, connect all the cities and spend the least amount.
Return the connects if can connect all the cities, otherwise return empty list.

Example

Gievn the connections = ["Acity","Bcity",1], ["Acity","Ccity",2], ["Bcity","Ccity",3]

Return ["Acity","Bcity",1], ["Acity","Ccity",2]

Notice

Return the connections sorted by the cost, or sorted city1 name if their cost is same, or sorted city2 if their city1 name is also same.

Solution:
Classic Kruskal algorithm, which is a greedy algorithm. The key idea is sort the connections by the cost. Then connect all the cities using Union-Find. Since we use the edges with the minimal costs first, and the Union-Find will result in a tree. After we connect all the edges, we will end with a minimum spanning tree. 

Code (Java):

/**
 * Definition for a Connection.
 * public class Connection {
 *   public String city1, city2;
 *   public int cost;
 *   public Connection(String city1, String city2, int cost) {
 *       this.city1 = city1;
 *       this.city2 = city2;
 *       this.cost = cost;
 *   }
 * }
 */
public class Solution {
    /**
     * @param connections given a list of connections include two cities and cost
     * @return a list of connections from results
     */
    public List<Connection> lowestCost(List<Connection> connections) {
        List<Connection> ans = new ArrayList<>();
        if (connections == null || connections.size() == 0) {
            return ans;
        }

        Collections.sort(connections, new MyConnectionComparator());

        int count = 0;
        Map<String, Integer> map = new HashMap<>();
        for (Connection connection : connections) {
            if (!map.containsKey(connection.city1)) {
                map.put(connection.city1, count++);
            }

            if (!map.containsKey(connection.city2)) {
                map.put(connection.city2, count++);
            }
        }

        int[] parents = new int[count];
        for (int i = 0; i < count; i++) {
            parents[i] = i;
        }

        for (Connection connection : connections) {
            int rootA = find(map.get(connection.city1), parents);
            int rootB = find(map.get(connection.city2), parents);

            if (rootA != rootB) {
                parents[rootA] = rootB;
                ans.add(connection);
            }
        }

        return ans.size() == count - 1 ? ans : new ArrayList<>();
    }

    private int find(int a, int[] parents) {
        int root = a;
        while (root != parents[root]) {
            root = parents[root];
        }

        while (root != a) {
            int temp = parents[a];
            parents[a] = root;
            a = temp;
        }

        return root;
    }
}

class MyConnectionComparator implements Comparator<Connection> {
    @Override
    public int compare(Connection a, Connection b) {
        if (a.cost != b.cost) {
            return a.cost - b.cost;
        }
        
        if (!a.city1.equals(b.city1)) {
            return a.city1.compareTo(b.city1);
        }
        
        return a.city2.compareTo(b.city2);
    }
}