Last Day Where You Can Still Cross LeetCode solution in Python

1
class UF:
2
    def __init__(self, m, n):
3
        self.n, self.loc_id, c_id = n, dict(), 0
4
        self.col_set = [set() for _ in range(m * n)]
5
        for i in range(m):  # Assign id for each location (i ,j)
6
            for j in range(n):
7
                self.loc_id[(i, j)] = c_id
8
                self.col_set[c_id].add(j)
9
                c_id += 1
10
        self.p = [i for i in range(m * n)]  # Initialize parents array `p`
11

12
    def find(self, i):
13
        if self.p[i] != i:
14
            self.p[i] = self.find(self.p[i])
15
        return self.p[i]
16

17
    def union(self, i, j):
18
        pi, pj = self.find(i), self.find(j)
19
        if pi != pj:
20
            self.p[pj] = pi  # Update `pi`
21
            self.col_set[pi] = (
22
                self.col_set[pi] | self.col_set[pj]
23
            )  # Take union of two sets (union all occupied columns)
24
        return (
25
            len(self.col_set[pi]) == self.n
26
        )  # if length of col_set[pi] == self.n, meaning this piece occupied all columns from 1 to `col` inclusive, meaning we are blocked
27

28

29
class Solution:
30
    def latestDayToCross(self, row: int, col: int, cells: List[List[int]]) -> int:
31
        uf, visited = UF(row, col), set()
32
        for i, (x, y) in enumerate(cells):
33
            x, y = x - 1, y - 1
34
            visited.add((x, y))
35
            for dx, dy in [
36
                (-1, -1),
37
                (-1, 0),
38
                (-1, 1),
39
                (1, -1),
40
                (1, 0),
41
                (1, 1),
42
                (0, 1),
43
                (0, -1),
44
            ]:
45
                _x, _y = x + dx, y + dy  # Check if neighbor is flooded
46
                if 0 <= _x < row and 0 <= _y < col and (_x, _y) in visited:
47
                    id1, id2 = uf.loc_id[(_x, _y)], uf.loc_id[(x, y)]
48
                    if uf.union(id1, id2):
49
                        return i  # Union two flooded piece and return index if union return True
50
        return -1