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

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

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.