LC 0317 [H] Shortest Distance from All Buildings - ALawliet/algorithms GitHub Wiki

the key parts are:

• having to build a solution matrix
• what it contains: [total travel distance, buildings reached] and we pop these off a min heap
• pruning step: matrix[nr][nc][1] == buildings_reached we only want to check buildings with the highest number reached so far

EMPTY = 0
BUILDING = 1
OBSTACLE = 2
DIRS = [(-1,0), (1,0), (0,1), (0,-1)]

class Solution(object):
    def shortestDistance(self, grid):
        n_rows = len(grid)
        n_cols = len(grid[0])
        
        # [total travel distance, buildings reached], use [] instead of () so we can reassign later
        matrix = [[ [0,0] for c in range(n_cols) ] for r in range(n_rows)]
        
        buildings_reached = 0    # count how many building we have visited
        
        def bfs(start):
            Q = deque([(start, 0)])
            while Q:
                node = Q.popleft()
                pos, distance = node
                r, c = pos
                for dr, dc in DIRS:
                    nr = r + dr
                    nc = c + dc
                    # only the position be visited by buildings_reached times will append to queue (pruning)
                    if 0 <= nr < n_rows and 0 <= nc < n_cols and matrix[nr][nc][1] == buildings_reached and grid[nr][nc] == EMPTY:
                        matrix[nr][nc][0] += distance + 1
                        matrix[nr][nc][1] = buildings_reached + 1
                        # matrix[nr][nc] = [matrix[nr][nc][0] + distance +1, buildings_reached + 1]
                        Q.append(([nr,nc], distance + 1))
        
        for r in range(n_rows):
            for c in range(n_cols):
                if grid[r][c] == BUILDING: # BFS from positions with a building
                    bfs([r,c])
                    buildings_reached += 1
                    
        # for r in range(n_rows):
            # print(matrix[r])

        res = float('inf')
        for r in range(n_rows):
            for c in range(n_cols):
                if matrix[r][c][1] == buildings_reached:
                    res = min(res, matrix[r][c][0])

        return res if res != float('inf') else -1