01 Matrix

Given an m x n binary matrix mat, return the distance of the nearest 0 for each cell.

The distance between two adjacent cells is 1.

Example 1:

Input: mat = [[0,0,0],[0,1,0],[0,0,0]]
Output: [[0,0,0],[0,1,0],[0,0,0]]

Example 2:

Input: mat = [[0,0,0],[0,1,0],[1,1,1]]
Output: [[0,0,0],[0,1,0],[1,2,1]]

Constraints:

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 104
• 1 <= m * n <= 104
• mat[i][j] is either 0 or 1.
• There is at least one 0 in mat.

SOLUTION:

import heapq
from collections import defaultdict

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        dist = [[0 if cell == 0 else float('inf') for cell in row] for row in mat]
        heap = []
        deleted = defaultdict(int)
        m = len(mat)
        n = len(mat[0])
        for i in range(m):
            for j in range(n):
                heapq.heappush(heap, (dist[i][j], i, j))
        while len(heap) > 0:
            currdist, i, j = heapq.heappop(heap)
            while (currdist, i, j) in deleted and len(heap) > 0:
                deleted[(currdist, i, j)] -= 1
                if deleted[(currdist, i, j)] == 0:
                    del deleted[(currdist, i, j)]
                currdist, i, j = heapq.heappop(heap)
            adj = []
            if i > 0:
                adj.append((i - 1, j))
            if i < m - 1:
                adj.append((i + 1, j))
            if j > 0:
                adj.append((i, j - 1))
            if j < n - 1:
                adj.append((i, j + 1))
            for a, b in adj:
                if dist[a][b] > 1 + currdist:
                    deleted[(dist[a][b], a, b)] += 1
                    dist[a][b] = 1 + currdist
                    heapq.heappush(heap, (1 + currdist, a, b))
        return dist