Recently, I've solved the "01 Matrix" LeetCode problem and the solution was accepted by the LeetCode OJ:
Given a matrix consists of 0 and 1, find the distance of the nearest 0 for each cell.
The distance between two adjacent cells is 1.
Input: 0 0 0 0 1 0 1 1 1 Output: 0 0 0 0 1 0 1 2 1
- The number of elements of the given matrix will not exceed 10,000.
- There are at least one 0 in the given matrix.
- The cells are adjacent in only four directions: up, down, left and right.
The idea behind the solution above is to use Dynamic Programming - starting with
0 cells work outwards putting not processed cells on the queue:
from collections import deque class Solution(object): def updateMatrix(self, matrix): """ :type matrix: List[List[int]] :rtype: List[List[int]] """ if not matrix: return matrix row_length = len(matrix) col_length = len(matrix) queue = deque() # put all 0 cells on queue, set all other cells to a big number for row_index in range(row_length): for col_index in range(col_length): if matrix[row_index][col_index] == 0: queue.append((row_index, col_index)) else: matrix[row_index][col_index] = 10000 # work from the 0 cells outwards while the queue is not empty while queue: row_index, col_index = queue.popleft() for i, j in [(row_index - 1, col_index), (row_index + 1, col_index), (row_index, col_index - 1), (row_index, col_index + 1)]: if 0 <= i < row_length and \ 0 <= j < col_length and \ matrix[i][j] > matrix[row_index][col_index] + 1: matrix[i][j] = matrix[row_index][col_index] + 1 queue.append((i, j)) return matrix
Even though the code works, I am not happy with the readability, in particular:
- setting the non-zero cells to "magical" 10000 does not look good
- getting the cell neighbors and checking if they are not out-of-bounds seems overly complicated
What would you improve code style and organization or time and space complexity wise?