Groups of Strings LeetCode solution in Python

1
class DSU:
2
    def __init__(self, n):
3
        self.root = list(range(n))
4

5
    def find(self, x):
6
        if self.root[x] != x:
7
            self.root[x] = self.find(self.root[x])
8
        return self.root[x]
9

10
    def union(self, x, y):
11
        self.root[self.find(x)] = self.find(y)
12

13

14
class Solution:
15
    def groupStrings(self, A: List[str]) -> List[int]:
16
        c = collections.defaultdict(int)
17
        for a in A:
18
            c["".join(sorted(a))] += 1
19

20
        A = list(set(["".join(sorted(a)) for a in A]))
21
        n = len(A)
22

23
        idx = collections.defaultdict(int)  # (binary representation -> index)
24
        dsu = DSU(n)  # dsu
25

26
        def add(base):
27
            for i in range(26):
28
                if not base & 1 << i:
29
                    yield base ^ 1 << i
30

31
        def dele(base):
32
            for i in range(26):
33
                if base & 1 << i:
34
                    if base - (1 << i) != 0:
35
                        yield base - (1 << i)
36

37
        def rep(base):
38
            pre, new = [], []
39
            for i in range(26):
40
                if base & 1 << i:
41
                    pre.append(i)
42
                else:
43
                    new.append(i)
44
            for p in pre:
45
                for n in new:
46
                    yield base - (1 << p) + (1 << n)
47

48
        for i, a in enumerate(A):
49
            base = 0
50
            for ch in a:
51
                base += 1 << ord(ch) - ord("a")
52
            idx[base] = i
53

54
        for base in idx.keys():
55
            for new in add(base):
56
                if new in idx:
57
                    dsu.union(idx[base], idx[new])
58
            for new in dele(base):
59
                if new in idx:
60
                    dsu.union(idx[base], idx[new])
61
            for new in rep(base):
62
                if new in idx:
63
                    dsu.union(idx[base], idx[new])
64

65
        group = collections.defaultdict(int)
66
        for a in A:
67
            base = 0
68
            for ch in a:
69
                base += 1 << ord(ch) - ord("a")
70
            cnum = c[a]
71
            group[dsu.find(idx[base])] += cnum
72

73
        return [len(group), max(group.values())]