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