1
class DSU(object):
2
    def __init__(self, N):
3
        self.par = list(range(N))
4
        self.rnk = [0] * N
5

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

11
    def union(self, x, y):
12
        xr, yr = self.find(x), self.find(y)
13
        if xr == yr:
14
            return False
15
        elif self.rnk[xr] < self.rnk[yr]:
16
            self.par[xr] = yr
17
        elif self.rnk[xr] > self.rnk[yr]:
18
            self.par[yr] = xr
19
        else:
20
            self.par[yr] = xr
21
            self.rnk[xr] += 1
22
        return True
23

24

25
class Solution:
26
    def swimInWater(self, grid):
27
        d, N = {}, len(grid)
28
        for i, j in product(range(N), range(N)):
29
            d[grid[i][j]] = (i, j)
30

31
        dsu = DSU(N * N)
32
        grid = [[0] * N for _ in range(N)]
33
        neib_list = [[0, 1], [0, -1], [1, 0], [-1, 0]]
34

35
        for i in range(N * N):
36
            x, y = d[i]
37
            grid[x][y] = 1
38
            for dx, dy in neib_list:
39
                if N > x + dx >= 0 and N > y + dy >= 0 and grid[x + dx][y + dy] == 1:
40
                    dsu.union((x + dx) * N + y + dy, x * N + y)
41

42
            if dsu.find(0) == dsu.find(N * N - 1):
43
                return i