542. 01 Matrix

542. 01 Matrix

Dijkstra to get shortest path from multiple source nodes. Using queue instead of heap (which makes it look like BFS) is because the distance increases by 1. It guarantees that nodes with smaller distance are at front of queue

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        heap = []
        distances = [[float("inf")] * n for _ in range(m)]
        for i in range(m):
            for j in range(n):
                if mat[i][j] == 0:
                    distances[i][j] = 0
                    heapq.heappush(heap, (0, i, j))
        
        while heap:
            distance, r, c = heapp.heappop(heap)
            for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                new_r, new_c = r + dr, c + dc
                if 0 <= new_r < m and 0 <= new_c < n and distances[new_r][new_c] > distance + 1:
                    distances[new_r][new_c] = distance + 1
                    queue.append((new_r, new_c))
        return distances