'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) -&gt; 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) -&gt; 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) -&gt; 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()
Repetition

I would write only one function to locate zeros:

def locate_zeros(matrix: IntegerMatrix) -&gt; 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) -&gt; Iterable[Position]:

def zero_out(matrix: IntegerMatrix) -&gt; 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.