/**
 * 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 class MyComp implements Comparator<Connection> {
        public int compare(Connection c1, Connection c2) {
            if (c1.cost != c2.cost) {
                return c1.cost - c2.cost;
            }
            if (c1.city1.equals(c2.city1) == false) {
                return c1.city1.compareTo(c2.city1);
            }
            if (c1.city2.equals(c2.city2) == false) {
                return c1.city2.compareTo(c2.city2);
            }
            return 0; // should not reach here
        }
    }

    public List<Connection> lowestCost(List<Connection> connections) {
        // Write your code here
        UnionFind uf = new UnionFind();
        Set<String> citySet = new HashSet<>();
        List<Connection> result = new ArrayList<>();
        PriorityQueue<Connection> conns = new PriorityQueue<>(new MyComp());
        for (Connection c: connections) {
            conns.offer(c);
            uf.add(c.city1);
            uf.add(c.city2);
            citySet.add(c.city1);
            citySet.add(c.city2);
        }

        while (true) {
            if (conns.size() == 0) break;

            Connection conn = conns.poll();
            if (uf.is_connected(conn.city1, conn.city2)) continue;
            result.add(conn);
            uf.union(conn.city1, conn.city2);
            // System.out.printf("(%d)%n", conn.cost);

        }

        // 圖論 - a tree will always have number_of_edges == (number_of_nodes - 1)
        if (result.size() != (citySet.size() - 1)) return new ArrayList<>();

        Collections.sort(result, new MyComp());

        return result;
    }

    class UnionFind {

        Map<String, String> city2parent = new HashMap<>();

        public void add(String c1) {
            if (city2parent.containsKey(c1)) return;
            city2parent.put(c1,c1); // be it's own parent as default
        }

        public String find_root(String c1) {
            if (city2parent.get(c1).equals(c1)) return c1;

            String root = find_root(city2parent.get(c1));
            city2parent.put(c1, root);

            return root;
        }

        public void merge(String c1, String c2) {
            String root2 = find_root(c2);
            city2parent.put(c1, root2);
        }

        public void union(String c1, String c2) {

            add(c1);
            add(c2);

            String root1 = find_root(c1);
            String root2 = find_root(c2);

            if (root1.equals(root2)) return;

            merge(root1, root2);  // attention: merge is at root level

        }

        public boolean is_connected(String c1, String c2){
            String root1 = find_root(c1);
            String root2 = find_root(c2);

            if (root1.equals(root2)) return true;

            return false;
        }
    }

}