Union Find Template for Matrix Data Structure

Union Find-Matrix.py

class Solution:
    def numIslands(self, mat: List[List[str]]) -> int:
        def find(a):
            if a != parent[a]:
                parent[a] = find(parent[a])

            return parent[a]
    
    
        def union(x,y):
            setX = find(x)
            setY = find(y)

            if setX != setY:
                if size[setX] >= size[setY]:
                    parent[setY] = setX
                    size[setX]+=size[setY]

                else:
                    parent[setX] = setY
                    size[setY]+=size[setX]

        directions = [(-1,0),(0,-1),(1,0),(0,1)]
        M, N = len(mat), len(mat[0])
        parent = {(i,j):(i,j) for i in range(M) for j in range(N)}
        size = {(i,j):1 for i in range(M) for j in range(N)}

        for i in range(M):
            for j in range(N):
                if mat[i][j] == "1":
                    for x,y in directions:
                        x+=i
                        y+=j

                        if x>=0 and x<M and y>=0 and y<N and mat[x][y] == "1":
                            union((i,j),(x,y))

        count = 0

        for component, par in parent.items():
            if component == par and mat[component[0]][component[1]] == "1":
                count+=1
                
        return count