Lintcode 573 · Build Post Office

Description

Given a 2D grid, each cell is either a wall 2, an house 1 or empty 0 (the number zero, one, two), find a place to build a post office so that the sum of the distance from the post office to all the houses is smallest.

Returns the sum of the minimum distances from all houses to the post office.Return -1 if it is not possible.
You cannot pass through wall and house, but can pass through empty.
You only build post office on an empty.

Note

• start from 1, and find all the reachable 0s. – BFS
• for each 0, record the number of 1s the 0 is linked with, and the distance to all linked 1s
• find the 0 which is connected with all 1s and has shortest distance.

Code

    def shortest_distance(self, grid: List[List[int]]) -> int:
        if not grid:
            return -1
        rows = len(grid)
        cols = len(grid[0])

        total_buildings = 0
        '''
        record the number of total buildings
        and the number of buildings each 0 can reach
        and their distance
        '''
        num_reach = [[0] * cols for _ in range(rows)]
        dist = [[0] * cols for _ in range(rows)]
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]

        for i in range(rows):
            for j in range(cols):
                if grid[i][j] == 1:
                '''
                for each house, bfs all the reachable 0s
                to get the dist and num_reach of each 0
                '''
                    total_buildings += 1
                    self.bfs(grid, i, j, dist, num_reach, directions)

        min_dist = float('inf')
        for i in range(rows):
            for j in range(cols):
                if num_reach[i][j] == total_buildings:
                    if dist[i][j] < min_dist:
                        min_dist = dist[i][j]

        if min_dist == float('inf'):
            return -1
        return min_dist
        
    def bfs(self, grid, x, y, dist, num_reach, directions):
        queue = collections.deque([(x, y)])
        visited = set([(x, y)])
        d = 0
        '''
        bfs by level for distance calculation
        bfs starts from 1 and go to find all the 0s that can be reached by this 1
        '''
        while queue:
        '''
        d is the distance from the start 1 to current level. 
        we are going to find the next 0s
        the distance to next 0 = current d + 1
        '''
            d += 1
            current_size = len(queue)
            for _ in range(current_size):
                current_pos = queue.popleft()
                for (n_x, n_y) in self.get_neighbors(current_pos, grid, directions):
                    if (n_x, n_y) in visited:
                        continue
                        '''
                        n_x, n_y is the position of next 0 that can be reached by current_pos
                        '''
                    dist[n_x][n_y] += d
                    num_reach[n_x][n_y] += 1
                    queue.append((n_x, n_y))
                    visited.add((n_x, n_y))
    
    def get_neighbors(self, pos, grid, directions):
        neighbors = []
        for (x, y) in directions:
            new_x, new_y = pos[0] + x, pos[1] + y
            if 0 <= new_x < len(grid) and 0 <= new_y < len(grid[0]) and grid[new_x][new_y] == 0:
                neighbors.append((new_x, new_y))
        return neighbors