Update Graphs/721_Accounts_Merge.java

Author: seanprashad

Graphs/721_Accounts_Merge.java

@@ -1,70 +1,62 @@
 class Solution {
-    public List<List<String>> accountsMerge(List<List<String>> accounts) {
-        UnionFind uf = new UnionFind(accounts.size());
-
-        Map<String, Integer> emailToAcc = new HashMap<>();
-
-        for (int i = 0; i < accounts.size(); i++) {
-            for (int j = 1; j < accounts.get(i).size(); j++) {
-                String email = accounts.get(i).get(j);
+    private class UnionFind {
+        private Map<String, String> parent;
 
-                if (emailToAcc.containsKey(email)) {
-                    int prevParent = emailToAcc.get(email);
-                    uf.union(prevParent, i);
-                } else {
-                    emailToAcc.put(email, i);
-                }
-            }
+        public UnionFind() {
+            parent = new HashMap<>();
         }
 
-        Map<Integer, Set<String>> emailsForAcc = new HashMap<>();
+        public String find(String x) {
+            // Upon adding a new item, its parent is itself
+            parent.putIfAbsent(x, x);
 
-        for (int i = 0; i < accounts.size(); i++) {
-            int parent = uf.find(i);
-            List<String> emails = accounts.get(i);
+            if (!parent.get(x).equals(x)) {
+                parent.put(x, find(parent.get(x))); // Path compression; directly specify the parent via recursion
+            }
 
-            emailsForAcc.putIfAbsent(parent, new HashSet<>());
-            emailsForAcc.get(parent).addAll(emails.subList(1, emails.size()));
+            return parent.get(x);
         }
 
-        List<List<String>> result = new ArrayList<>();
+        public void union(String x, String y) {
+            String rootX = find(x);
+            String rootY = find(y);
 
-        for (int key : emailsForAcc.keySet()) {
-            List<String> temp = new ArrayList<>();
-            temp.addAll(emailsForAcc.get(key));
-            Collections.sort(temp);
-            temp.add(0, accounts.get(key).get(0));
-            result.add(temp);
+            if (!rootX.equals(rootY)) {
+                parent.put(rootX, rootY);
+            }
         }
-
-        return result;
     }
 
-    private class UnionFind {
-        private int[] parents;
+    public List<List<String>> accountsMerge(List<List<String>> accounts) {
+        Map<String, String> emailToName = new HashMap<>();
+        UnionFind uf = new UnionFind();
 
-        public UnionFind(int n) {
-            parents = new int[n];
+        for (List<String> account : accounts) {
+            String name = account.get(0);
+            String initialEmail = account.get(1);
 
-            for (int i = 0; i < parents.length; i++) {
-                parents[i] = i;
+            for (int i = 1; i < account.size(); i++) {
+                String email = account.get(i);
+                emailToName.putIfAbsent(email, name);
+                uf.union(initialEmail, email);
             }
         }
 
-        public int find(int node) {
-            if (parents[node] == node) {
-                return node;
-            }
-            parents[node] = find(parents[parents[node]]);
-            return parents[node];
+        Map<String, TreeSet<String>> unions = new HashMap<>();
+        for (String email : emailToName.keySet()) {
+            String root = uf.find(email);
+            unions.computeIfAbsent(root, x -> new TreeSet<>()).add(email);
         }
 
-        public void union(int i, int j) {
-            int p1 = find(i), p2 = find(j);
-            if (p1 == p2) {
-                return;
-            }
-            parents[p2] = p1;
+        List<List<String>> result = new ArrayList<>();
+
+        for (Map.Entry<String, TreeSet<String>> entry : unions.entrySet()) {
+            List<String> temp = new ArrayList<>();
+            temp.add(emailToName.get(entry.getKey()));
+            temp.addAll(entry.getValue());
+            result.add(temp);
         }
+
+        return result;
     }
-}
+}