Find All People With Secret LeetCode solution in Python

1
# Runtime: 6190 ms (Top 5.05%) | Memory: 57.3 MB (Top 48.67%)
2
class Solution:
3
    def findAllPeople(
4
        self, n: int, meetings: List[List[int]], firstPerson: int
5
    ) -> List[int]:
6

7
        class UnionFind:
8
            def __init__(self):
9
                self.parents = {}
10
                self.ranks = {}
11

12
            def insert(self, x):
13
                if x not in self.parents:
14
                    self.parents[x] = x
15
                    self.ranks[x] = 0
16

17
            def find_parent(self, x):
18
                if self.parents[x] != x:
19
                    self.parents[x] = self.find_parent(self.parents[x])
20
                return self.parents[x]
21

22
            def union(self, x, y):
23
                self.insert(x)
24
                self.insert(y)
25
                x, y = self.find_parent(x), self.find_parent(y)
26
                if x == y:
27
                    return
28
                if self.ranks[x] > self.ranks[y]:
29
                    self.parents[y] = x
30
                else:
31
                    self.parents[x] = y
32
                    if self.ranks[x] == self.ranks[y]:
33
                        self.ranks[y] += 1
34

35
        time2meets = defaultdict(list)
36
        for x, y, t in meetings:
37
            time2meets[t].append((x, y))
38
        time2meets = sorted(time2meets.items())
39

40
        curr_know = set([0, firstPerson])
41

42
        for time, meets in time2meets:
43
            uf = UnionFind()
44
            for x, y in meets:
45
                uf.union(x, y)
46

47
            groups = defaultdict(set)
48
            for idx in uf.parents:
49
                groups[uf.find_parent(idx)].add(idx)
50

51
            for group in groups.values():
52
                if group & curr_know:
53
                    curr_know.update(group)
54

55
        return list(curr_know)