[LeetCode] 721. Accounts Merge

 Problem : https://leetcode.com/problems/accounts-merge/

Use union-find to group email address of the same account.  Remember to sort the mail list in the end.

class UnionFind {
    int[] parent;
    
    UnionFind(int size) {
        parent = new int[size];
        
        for (int i = 0; i < size; i++) {
            parent[i] = i;
        }
    }
    
    void union(int a, int b) {
        int ap = find(a);
        int bp = find(b);
        
        if (ap != bp) {
            parent[ap] = bp;
        }
    }
    
    int find(int a) {
        if (parent[a] == a) {
            return a;
        }
        
        parent[a] = find(parent[a]);
        
        return parent[a];
    }
}


class Solution {
    public List<List<String>> accountsMerge(List<List<String>> accounts) {
        int ID = 0;
        
        // convert mail to ids
        HashMap<String, Integer> mailIds = new HashMap();
        HashMap<Integer, String> idToMail = new HashMap();
        HashMap<String, Integer> mailToName = new HashMap();
        
        for (int i = 0; i < accounts.size(); i++) {
            for (int j = 1; j < accounts.get(i).size(); j++) {
                if (mailIds.containsKey(accounts.get(i).get(j))) {
                    continue;
                }
                
                mailIds.put(accounts.get(i).get(j), ID);
                idToMail.put(ID, accounts.get(i).get(j));
                ID++;
                mailToName.put(accounts.get(i).get(j), i);
            }    
        }
        
        // use union-find to group mail ids
        UnionFind uf = new UnionFind(ID);
        
        for (int i = 0; i < accounts.size(); i++) {
            for (int j = 2; j < accounts.get(i).size(); j++) {
                uf.union(mailIds.get(accounts.get(i).get(j-1)), mailIds.get(accounts.get(i).get(j)));        
            }
        }
        
        // construct the mail list
        HashMap<Integer, List<String>> groups = new HashMap();
        
        for (int i = 0; i < ID; i++) {
            String email = idToMail.get(i);
       
            if (groups.containsKey(uf.find(i))) {
                groups.get(uf.find(i)).add(email);
            } else {
                List<String> emailList = new ArrayList();
                emailList.add(email);
                groups.put(uf.find(i), emailList);
            }
        }
        
        List<List<String>> result = new ArrayList();
        for (int pid : groups.keySet()) {
            List<String> list = new ArrayList();
            
            // sort mails
            Collections.sort(groups.get(pid));
            String name = accounts.get(mailToName.get(idToMail.get(pid))).get(0);
            
            list.add(name);
            list.addAll(groups.get(pid));
            result.add(list);
        }
        
        return result;
    }
}