# 'Zero out' a matrix using Python

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()

• Is numpy not allowed? Sep 19, 2016 at 18:34
• The problem does not say you can not use numpy. I am solving these problems to learn programming and to prepare for technical interviews. Sep 19, 2016 at 19:11

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

• Why do we need two for rowi loops? Sep 19, 2016 at 19:20
• @newToProgramming: typo, sorry, fixed Sep 19, 2016 at 19:36
• ~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.
– DSM
Sep 3, 2018 at 14:52
• @DSM Good catch, fixed Sep 6, 2018 at 15:56

### 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.

• 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. Sep 18, 2016 at 22:47
• What is zeros_positions?
– vnp
Sep 18, 2016 at 22:55
• @vnp locate_zeros in the matrix Sep 19, 2016 at 12:34