Union find clean python code


  • 0
    class UnionFind(object):
        def __init__(self, board):
            self.m = len(board)
            self.n = len(board[0])
            self.length = self.m * self.n
    
            self.id = [None] * self.length
            self.size = [1] * self.length
            self.surrounded = [True] * self.length
    
            [operator.setitem(self.id, *([self.genIndex(i, j)] * 2))
             for i, row in enumerate(board)
             for j, val in enumerate(row) if val == 'O']
    
        def genIndex(self, i, j):
            return self.n * i + 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.surrounded[more] = self.surrounded[less] and self.surrounded[more]
    
    
    class Solution(object):
        def solve(self, board):
            if not board:
                return
    
            uf = UnionFind(board)
    
            for i, row in enumerate(board):
                for j, val in enumerate(row):
                    if val != 'O':
                        continue
    
                    index = uf.genIndex(i, j)
    
                    [uf.union(index, uf.genIndex(y, z))
                     for x, y, z in ((i, i - 1, j), (j, i, j - 1))
                     if x > 0 and board[y][z] == 'O']
    
                    if i == 0 or j == 0 or i == uf.m - 1 or j == uf.n - 1:
                        uf.surrounded[uf.find(index)] = False
    
            [operator.setitem(board[i], j, 'X')
             for i in xrange(uf.m)
             for j in xrange(uf.n)
             if board[i][j] == 'O' and uf.surrounded[uf.find(uf.genIndex(i, j))]]

Log in to reply
 

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