Making A Large Island LeetCode solution in Python

1
class Solution:
2
    def largestIsland(self, grid: List[List[int]]) -> int:
3
        m, n = len(grid), len(grid[0])
4
        # parent array to keey track
5
        parent = list(range(m * n))
6
        # rank array used for union by rank and size calculation
7
        rank = [1 for _ in range(m * n)]
8

9
        # standard DSU find function
10
        def find(x):
11
            while x != parent[x]:
12
                parent[x] = parent[parent[x]]
13
                x = parent[x]
14
            return x
15

16
        # standard DSU union by rank function
17
        def union(x, y):
18
            parX, parY = find(x), find(y)
19
            if parX == parY:
20
                return
21
            if rank[parX] >= rank[parY]:
22
                rank[parX] += rank[parY]
23
                parent[parY] = parX
24
            else:
25
                rank[parY] += rank[parX]
26
                parent[parX] = parY
27

28
        # Step 1: join the land
29

30
        # for each island we perform union operation
31
        for i, row in enumerate(grid):
32
            for j, val in enumerate(row):
33

34
                # important condition: if we have a water body, we set its rank to 0 (will be used in the next step)
35
                if not val:
36
                    rank[i * m + j] = 0
37
                    continue
38

39
                # performing union of land bodies
40
                for x, y in [(i - 1, j), (i + 1, j), (i, j + 1), (i, j - 1)]:
41
                    # outside of grid check
42
                    if not (0 <= x < m and 0 <= y < n):
43
                        continue
44

45
                    if grid[x][y]:
46
                        union(i * m + j, x * m + y)
47

48
        # Step 2: convert a water body (if present)
49

50
        # the minimum final ans will always be the size of the largest land present
51
        ans = max(rank)
52
        for i, row in enumerate(grid):
53
            for j, val in enumerate(row):
54

55
                # we dont need to do anything if we encounter a land
56
                if val:
57
                    continue
58

59
                neighbours = set()
60
                res = 0
61

62
                #
63
                for x, y in [(i - 1, j), (i + 1, j), (i, j + 1), (i, j - 1)]:
64
                    # outside of grid check
65
                    if not (0 <= x < m and 0 <= y < n):
66
                        continue
67

68
                    # checking unique neighbours by adding the parent to the set
69
                    # here we dont care if the neighbour is water as its rank is 0 so it contributes nothing
70
                    idx = x * m + y
71
                    neighbours.add(find(idx))
72

73
                # Once we have all unique neighbours, just add their ranks
74
                for idx in neighbours:
75
                    res += rank[idx]
76

77
                # res + 1 because we convert the current cell (i,j) to land too
78
                ans = max(ans, res + 1)
79

80
        return ans