5
\$\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\$

1 Answer 1

5
\$\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\$

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.