Accounts Merge LeetCode solution in Java

1
class Solution {
2
  class UnionFind {
3
    int[] parent;
4
    int[] weight;
5

6
    public UnionFind(int num) {
7
      parent = new int[num];
8
      weight = new int[num];
9

10
      for (int i = 0; i < num; i++) {
11
        parent[i] = i;
12
        weight[i] = 1;
13
      }
14
    }
15

16
    public void union(int a, int b) {
17
      int rootA = root(a);
18
      int rootB = root(b);
19

20
      if (rootA == rootB) {
21
        return;
22
      }
23

24
      if (weight[rootA] > weight[rootB]) {
25
        parent[rootB] = rootA;
26
        weight[rootA] += weight[rootB];
27
      } else {
28
        parent[rootA] = rootB;
29
        weight[rootB] += weight[rootA];
30
      }
31
    }
32

33
    public int root(int a) {
34
      if (parent[a] == a) {
35
        return a;
36
      }
37

38
      parent[a] = root(parent[a]);
39
      return parent[a];
40
    }
41
  }
42

43
  public List<List<String>> accountsMerge(List<List<String>> accounts) {
44
    int size = accounts.size();
45

46
    UnionFind uf = new UnionFind(size);
47

48
    // prepare a hash with unique email address as key and index in accouts as value
49
    HashMap<String, Integer> emailToId = new HashMap<>();
50
    for (int i = 0; i < size; i++) {
51
      List<String> details = accounts.get(i);
52
      for (int j = 1; j < details.size(); j++) {
53
        String email = details.get(j);
54

55
        // if we have already seen this email before, merge the account  "i" with previous account
56
        // else add it to hash
57
        if (emailToId.containsKey(email)) {
58
          uf.union(i, emailToId.get(email));
59
        } else {
60
          emailToId.put(email, i);
61
        }
62
      }
63
    }
64

65
    // prepare a hash with index in accounts as key and list of unique email address for that
66
    // account as value
67
    HashMap<Integer, List<String>> idToEmails = new HashMap<>();
68
    for (String key : emailToId.keySet()) {
69
      int root = uf.root(emailToId.get(key));
70

71
      if (!idToEmails.containsKey(root)) {
72
        idToEmails.put(root, new ArrayList<String>());
73
      }
74

75
      idToEmails.get(root).add(key);
76
    }
77

78
    // collect the emails from idToEmails, sort it and add account name at index 0 to get the final
79
    // list to add to final return List
80
    List<List<String>> mergedDetails = new ArrayList<>();
81
    for (Integer id : idToEmails.keySet()) {
82
      List<String> emails = idToEmails.get(id);
83
      Collections.sort(emails);
84
      emails.add(0, accounts.get(id).get(0));
85

86
      mergedDetails.add(emails);
87
    }
88

89
    return mergedDetails;
90
  }
91
}