'Zero out' a matrix using Python

Question

Below is my code for an attempted solution to cracking the coding interview exercise 1.8 written in python 3.5. The problem statement is:

Write an algorithm such that if an element in an MxN matrix is 0, its entire row and column are set to 0.

I looked at the hints for this problem so I do know that my solution is not the most space efficient possible. I believe the time complexity of the code below is \$O(M*N)\$ and the space complexity is \$O(M + N)\$. I wrote small unit tests to test my code. Looking for feedback on what needs to be improved, especially in terms of readability.

import unittest
def locate_zero_rows(matrix: list) -> list:
 """Given an NxM matrix find the rows that contain a zero."""
 zero_rows = [False for _ in range(len(matrix))]
 for row_num, row in enumerate(matrix):
 for col_num, element in enumerate(row):
 if element == 0:
 zero_rows[row_num] = True
 return zero_rows
def locate_zero_cols(matrix: list) -> list:
 """Given an NxM matrix find the columns that contain a zero."""
 zero_cols = [False for _ in range(len(matrix[0]))]
 for row_num, row in enumerate(matrix):
 for col_num, element in enumerate(row):
 if element == 0:
 zero_cols[col_num] = True
 return zero_cols
def zero_out(matrix: list) -> list:
 """Given an NxM matrix zero out all rows and columns that contain at least one zero."""
 zero_rows, zero_cols = locate_zero_rows(matrix), locate_zero_cols(matrix)
 for row_num, row in enumerate(matrix):
 for col_num, element in enumerate(row):
 if zero_rows[row_num] or zero_cols[col_num]:
 matrix[row_num][col_num] = 0
 return matrix
class MyTest(unittest.TestCase):
 def test_locate_zero_rows(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zero_rows = [False, True, True]
 self.assertSequenceEqual(locate_zero_rows(matrix), zero_rows)
 def test_locate_zero_cols(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zero_cols = [True, True, False, True]
 self.assertSequenceEqual(locate_zero_cols(matrix), zero_cols)
 def test_zero_out(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zeroed_out_matrix = [[0, 0, 2, 0],
 [0, 0, 0, 0],
 [0, 0, 0, 0]]
 self.assertSequenceEqual(zero_out(matrix), zeroed_out_matrix)
if __name__ == '__main__':
 unittest.main()

• TheBlackCat: Is numpy not allowed?
• Average: The problem does not say you can not use numpy. I am solving these problems to learn programming and to prepare for technical interviews.

Answer 1

The big issue I can see is that you search every element. You don't need to do this, you only need to check if 0 is anywhere in that row or column. You can use the in operation to test this. It will short-circuit after the first 0 is found, avoiding having to search the entire row or column. This won't reduce the time complexity, but will improve the best-case performance considerably.

Second, you can use zip to switch between rows and columns.

Third, you can reduce that check to a simple list comprehension, generator expression, or generator function.

Fourth, you can use ~all to detect if there are any zeros in a given sequence. This is slightly faster than in.

Finally, since you are making the changes in-place, you don't need to return the modified matrix.

So here is my version:

import unittest
def locate_zero_rows(matrix: list) -> list:
 """Given an NxM matrix find the rows that contain a zero."""
 return [i for i, row in enumerate(matrix) if not all(row)]
def locate_zero_cols(matrix: list) -> list:
 """Given an NxM matrix find the columns that contain a zero."""
 return locate_zero_rows(zip(*matrix))
def zero_out(matrix: list) -> None:
 """Given an NxM matrix zero out all rows and columns that contain at least one zero."""
 zero_rows = locate_zero_rows(matrix)
 zero_cols = locate_zero_cols(matrix)
 ncol = len(matrix[0])
 for rowi in zero_rows:
 matrix[rowi] = [0]*ncol
 for coli in zero_cols:
 for row in matrix:
 row[coli] = 0
class MyTest(unittest.TestCase):
 def test_locate_zero_rows(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zero_rows = [1, 2]
 self.assertSequenceEqual(locate_zero_rows(matrix), zero_rows)
 def test_locate_zero_cols(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zero_cols = [0, 1, 3]
 self.assertSequenceEqual(locate_zero_cols(matrix), zero_cols)
 def test_zero_out(self):
 matrix = [[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]]
 zeroed_out_matrix = [[0, 0, 2, 0],
 [0, 0, 0, 0],
 [0, 0, 0, 0]]
 zero_out(matrix)
 self.assertSequenceEqual(matrix, zeroed_out_matrix)
if __name__ == '__main__':
 unittest.main()

You can improve this further by making the column list comprehension a expression. I think this will give this a O(M) space complexity:

def zero_out(matrix: list) -> None:
 """Given an NxM matrix zero out all rows and columns that contain at least one zero."""
 zero_cols = (i for i, col in enumerate(zip(*matrix)) if not all(col))
 zero_rows = [i for i, row in enumerate(matrix) if not all(row)]
 ncol = len(matrix[0])
 for coli in zero_cols:
 for row in matrix:
 row[coli] = 0
 for rowi in zero_rows:
 matrix[rowi] = [0]*ncol

You can't make both comprehensions with this structure because changes to one would be reflected in the other.

It is possible to make both comprehensions using itertools.zip_longest, but you don't gain any space complexity (at least for matrices where N and M are similar), and it hurts your performance.

If you can use numpy, this can be simplified enormously:

import numpy as np
import unittest
def zero_out(matrix: np.array) -> None:
 """Given an NxM matrix zero out all rows and columns that contain at least one zero."""
 zero_cols = ~matrix.all(axis=0)
 zero_rows = ~matrix.all(axis=1)
 matrix[:, zero_cols] = 0
 matrix[zero_rows, :] = 0
class MyTest(unittest.TestCase):
 def test_zero_out(self):
 matrix = np.array([[5, 3, 2, 1],
 [-3, 0, 5, 0],
 [0, -1, 2, 6]])
 zeroed_out_matrix = np.array([[0, 0, 2, 0],
 [0, 0, 0, 0],
 [0, 0, 0, 0]])
 zero_out(matrix)
 np.testing.assert_array_equal(matrix, zeroed_out_matrix)
if __name__ == '__main__':
 unittest.main()

Edit: added ~all. Edit 2: Add numpy Edit 3: use not all instead of ~all

• Average: Why do we need two for rowi loops?
• TheBlackCat: @newToProgramming: typo, sorry, fixed
• DSM: ~all(row) isn't a good idea. ~ is a bitwise not, not a logical not, and since True == 1, ~True == -2, meaning that bool(True) == bool(~True) == True.
• TheBlackCat: @DSM Good catch, fixed

Answer 2

Repetition

I would write only one function to locate zeros:

def locate_zeros(matrix: IntegerMatrix) -> Iterable[Position]:
 """Given an NxM matrix find the positions that contain a zero."""
 for row_num, row in enumerate(matrix):
 for col_num, element in enumerate(row):
 if element == 0:
 yield (col_num, row_num)

And use it in zero_out like this:

 if row_num in (x[1] for x in zeros_positions) or col_num in (x[0] for x in zeros_positions):
 matrix[row_num][col_num] = 0

Type hints

Given that you specifically mentioned Python 3.5 and that you already have something like type hints on your functions, I suggest you go all the way with mypy compatible type hints.

from typing import List, Any, Iterable, Tuple
Position = Tuple(int, int)
IntegerMatrix = List[List[int]]
def locate_zeros(matrix: IntegerMatrix) -> Iterable[Position]:
def zero_out(matrix: IntegerMatrix) -> IntegerMatrix:

This way you can statically check your code has the correct types like in natively statically types languages and give the user much more detailed information on the types.

• Average: thank you. I will try and understand the space and time complexity of your implementation of locate_zeros, I assume it is O(1) space O(m*n) time? Though in retrospect i should have combined the two functions into one even without implementing it as a generator I assume.
• vnp: What is zeros_positions?
• Caridorc: @vnp locate_zeros in the matrix