blah.py

class Solution:
    def updateMatrix(self, matrix: List[List[int]]) -> List[List[int]]:
        
        dimX, dimY = len(matrix), len(matrix[0])
        
        for i in range(dimX+dimY):
            #print("mat", matrix)
            for x in range(dimX):
                for y in range(dimY):
                    
                    if matrix[x][y] == 0:
                        continue
                    
                    ofX, ofY = [0,0,-1,1], [-1,1,0,0]
                    
                    neighbours = [matrix[x+ofX[j]][y+ofY[j]] for j in range(4) if 
                                  x+ofX[j] >= 0 and x+ofX[j] < dimX and
                                  y+ofY[j] >= 0 and y+ofY[j] < dimY]
                    
                    #print(neighbours)
                    matrix[x][y] = 1+min(neighbours)
                    
        return matrix