AccountMerge.java

class Solution {
    public List<List<String>> accountsMerge(List<List<String>> accounts) {
         List<List<String>> results = new ArrayList<>();
        if (accounts == null || accounts.size() < 1) {
            return results;
        }
        
        HashMap<String, List<Integer>> emailToIds = mapEmailToIds(accounts);
        UnionFind unionFind = linkAccounts(accounts, emailToIds);
        HashMap<Integer, HashSet<String>> idToEmails = mapIdToEmails(accounts, unionFind);
        return generateAccountInfos(accounts, idToEmails);
    }
    
    private UnionFind linkAccounts(List<List<String>> accounts, HashMap<String, List<Integer>> emailToIds) {
        UnionFind unionFind = new UnionFind(accounts.size());
        for (String email : emailToIds.keySet()) {
            List<Integer> ids = emailToIds.get(email);
            
            for (int i = 1; i < ids.size(); i++) {
                unionFind.union(ids.get(i - 1), ids.get(i));
            }
        }
        return unionFind;
    }
    
    private HashMap<String, List<Integer>> mapEmailToIds(List<List<String>> accounts) {
        HashMap<String, List<Integer>> emailToIds = new HashMap<>();
        
        for (int id = 0; id < accounts.size(); id++) {
            List<String> infos = accounts.get(id);
            
            for (int i = 1; i < infos.size(); i++) {
                String email = infos.get(i);
                List<Integer> ids = emailToIds.getOrDefault(email, new ArrayList<>());
                ids.add(id);
                emailToIds.put(email, ids);
            }
        }
        return emailToIds;
    }
    
    private HashMap<Integer, HashSet<String>> mapIdToEmails(List<List<String>> accounts, UnionFind unionFind) {
        HashMap<Integer, HashSet<String>> idToEmails = new HashMap<>();
        for (int id = 0; id < accounts.size(); id++) {
            int rootId = unionFind.find(id);
            
            HashSet<String> uniqueEmails = idToEmails.getOrDefault(rootId, new HashSet<>());
            List<String> emails = accounts.get(id);
            for (int i = 1; i < emails.size(); i++) {
                uniqueEmails.add(emails.get(i));
            }
            idToEmails.put(rootId, uniqueEmails);
        }
        return idToEmails;
    }
    
    private List<List<String>> generateAccountInfos(List<List<String>> accounts, HashMap<Integer, HashSet<String>> idToEmails) {
        List<List<String>> results = new ArrayList<>();
        
        for (int id : idToEmails.keySet()) {
            String name = accounts.get(id).get(0);
            
            List<String> emails = new ArrayList<>(idToEmails.get(id));
            Collections.sort(emails);
            
            List<String> infos = new ArrayList<>();
            infos.add(name);
            
            for (String email: emails) {
                infos.add(email);
            }
            results.add(infos);
        }
        return results;
    }
    
    private static final class UnionFind {
        private int[] roots;
        
        UnionFind(int size) {
            roots = new int[size];
            for (int i = 0; i < size; i++) {
                roots[i] = i;
            }
        }
        
        void union(int node1, int node2) {
            int root1 = find(node1);
            int root2 = find(node2);
            if (root1 == root2) {
                return;
            }
            roots[root1] = root2;
        }
        
        int find(int node) {
            int root = roots[node];
            if (root == node) {
                return root;
            }
            return find(root);
        }
    }
}