Problem statement

https://leetcode.com/problems/01-matrix/

Solution

In this problem for each cell we need to find the closest distance to cell with 0 inside. We can reverse this problem: we start with all cells with 0 inside and run bfs - we can call it multi-source bfs, because in the first moment of time we start from several different cells. Also here I use variant of bfs when we go level by level: first we found all cells with distance 1, then with distance 2 and so on - you can do in in classical way as well. When we traverse some cell we update its value, so in the end we just return matrix M.

Complexity

Time complexity is O(mn), space complexity as well.

Code

class Solution:
    def updateMatrix(self, M):
        m, n, queue, visited = len(M), len(M[0]), deque(), set()
        for y, x in product(range(m), range(n)):
            if M[y][x] == 0: 
                queue.append((y, x))
                visited.add((y, x))
                
        dirs = [[1,0],[-1,0],[0,1],[0,-1]]
        level = 0
        
        while queue:
            level += 1
            for _ in range(len(queue)):
                y, x = queue.popleft()
                for dy, dx in dirs:
                    if 0<=y+dy<m and 0<=x+dx<n and (y+dy, x+dx) not in visited:
                        M[y+dy][x+dx] = level
                        queue.append((y+dy, x+dx))
                        visited |= {(y+dy, x+dx)}
        
        return M