Lintcode 629. Minimum Spanning Tree

/**
 * 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 {

    private class DisjointSet<T> {
        private Map<T, T> parents = new HashMap<T, T>();
        private Map<T, Integer> sizes = new HashMap<T, Integer>();

        public DisjointSet() {}

        public void add(T e) {
            if(parents.containsKey(e)) {
                return;
            }

            parents.put(e, e);
            sizes.put(e, 1);
        }

        public T find(T e) {
            if(!e.equals(parents.get(e))) {
                parents.put(e, find(parents.get(e)));
            }
            return parents.get(e);
        }

        public boolean union(T e1, T e2) {
            T r1 = find(e1);
            T r2 = find(e2);

            if(r1.equals(r2)) {
                return false;
            }

            int size1 = sizes.get(r1);
            int size2 = sizes.get(r2);
            if(size1 > size2) {
                parents.put(r2, r1);
                sizes.put(r1, size1 + size2);
            } else {
                parents.put(r1, r2);
                sizes.put(r2, size1 + size2);
            }
            return true;
        }

        public int size() {
            return parents.size();
        }
    }

    /**
     * @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) {
        Collections.sort(connections, (Connection c1, Connection c2) -> {
            if(c1.cost != c2.cost) {
                return c1.cost - c2.cost;
            } else if(!c1.city1.equals(c2.city1)) {
                return c1.city1.compareTo(c2.city1);
            } else {
                return c1.city2.compareTo(c2.city2);
            }
        });

        List<Connection> mst = new ArrayList<Connection>();

        DisjointSet<String> ds = new DisjointSet<String>();
        for(Connection connection : connections) {
            String city1 = connection.city1;
            String city2 = connection.city2;

            ds.add(city1);
            ds.add(city2);

            if(ds.union(city1, city2)) {
                mst.add(connection);
            }
        }

        // if mst doesn't form a spanning tree, return an empty list
        if(mst.size() != ds.size() - 1) {
            return new ArrayList(0);
        }
        return mst;
    }
}