Groups of Strings LeetCode solution in C++

1
// OJ: https://leetcode.com/problems/groups-of-strings/
2
// Author: github.com/lzl124631x
3
// Time: O(26 * N * alpha(N))
4
// Space: O(26 * N)
5
class UnionFind {
6
  vector<int> id, rank, size;
7
  int cnt;
8

9
 public:
10
  UnionFind(int n) : id(n), rank(n, 0), size(n, 1), cnt(n) {
11
    iota(begin(id), end(id), 0);
12
  }
13
  int find(int a) {
14
    return id[a] == a ? a : (id[a] = find(id[a]));
15
  }
16
  void connect(int a, int b) {
17
    int x = find(a), y = find(b);
18
    if (x == y) return;
19
    if (rank[x] > rank[y]) {
20
      id[y] = x;
21
      size[x] += size[y];
22
    } else {
23
      id[x] = y;
24
      size[y] += size[x];
25
      if (rank[x] == rank[y]) rank[y]++;
26
    }
27
    --cnt;
28
  }
29
  int getSize(int a) {
30
    return size[find(a)];
31
  }
32
  int getCount() {
33
    return cnt;
34
  }
35
};
36
class Solution {
37
 public:
38
  vector<int> groupStrings(vector<string> &A) {
39
    int N = A.size();
40
    UnionFind uf(N);
41
    unordered_map<int, int> m, delMap;
42
    m.reserve(N);
43
    for (int i = 0; i < N; ++i) {
44
      int h = 0;
45
      for (char c : A[i]) h |= 1 << (c - 'a'); // `h` is the bitmask representation of `A[i]`
46
      for (int j = 0; j < 26; ++j) {
47
        if (h >> j & 1) {         // if `h`'s j-th bit is 1
48
          int del = h ^ (1 << j); // `del` is the bitmask after deleting the `j`-th bit
49
          if (m.count(del)) uf.connect(i, m[del]); // Connect `A[i]` with its deletion variant
50
          if (delMap.count(del))
51
            uf.connect(i, delMap[del]);
52
          else
53
            delMap[del] = i;
54
        } else {
55
          int add = h | (1 << j); // `add` is the bitmask after adding `j`-th bit
56
          if (m.count(add)) uf.connect(i, m[add]);
57
        }
58
      }
59
      m[h] = i;
60
    }
61
    int mx = 1;
62
    for (int i = 0; i < N; ++i) mx = max(mx, uf.getSize(i));
63
    return {uf.getCount(), mx};
64
  }
65
};