Clean union find python code


  • 0
    class UnionFind(object):
        def __init__(self, grid):
            self.cols = len(grid[0])
            self.size = [1] * (len(grid) * self.cols)
            self.id = [
                self.genIndex(i, j) if num == '1' else None
                for i, row in enumerate(grid)
                for j, num in enumerate(row)
            ]
            self.count = len(filter(lambda x: x is not None, self.id))
    
        def genIndex(self, i, j):
            return i * self.cols + j
    
        def find(self, p):
            while p != self.id[p]:
                self.id[p] = self.id[self.id[p]]
                p = self.id[p]
            return p
    
        def union(self, p, q):
            idp, idq = map(self.find, (p, q))
            if idp == idq:
                return
    
            less, more = (
                (idp, idq) if self.size[idp] < self.size[idq] else (idq, idp))
    
            self.id[less] = self.id[more]
            self.size[more] += self.size[less]
    
            self.count -= 1
    
    
    class Solution(object):
        def numIslands(self, grid):
            if not grid:
                return 0
    
            uf = UnionFind(grid)
    
            [uf.union(uf.genIndex(i, j), uf.genIndex(y, z))
             for i, row in enumerate(grid)
             for j, num in enumerate(row)
             for x, y, z in [(i, i - 1, j), (j, i, j - 1)]
             if num == '1' and x > 0 and grid[y][z] == '1']
    
            return uf.count

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.