3
\$\begingroup\$

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.

Example:

Input:
0 0 0
0 1 0
1 1 1

Output:
0 0 0
0 1 0
1 2 1

Note:

  • 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[0])

        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?

\$\endgroup\$
4
\$\begingroup\$

By storing the return value in a different variable and placing the distance in the queue, the magic number can be avoided:

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[0])

        result = [[None for j in range(col_length)] for i in range(row_length)]

        queue = deque()

        # put all 0 cells in queue, set all other cells to a big number
        for row_index in range(row_length):
            for col_index in range(col_length):
                if 0 == matrix[row_index][col_index]:
                    queue.append((row_index, col_index, 0))

        # work from the 0 cells outwards while the queue is not empty
        while queue:
            i, j, dist = queue.popleft()
            if 0 <= i < row_length and 0 <= j < col_length and result[i][j] is None:
                result[i][j] = dist
                queue.append((i-1, j, dist+1))
                queue.append((i+1, j, dist+1))
                queue.append((i, j-1, dist+1))
                queue.append((i, j+1, dist+1))

        return result
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.