- find cols and rows that have a zero
- set value of coords with those cols or rows to zero

# [1.8] Zero Matrix: Write an algorithm such that if an element
# in an MxN matrix is 0, its entire row and columns are set to 0.
#Time Complexity: O(MN)
#Space Complexity: O(M+N)
import unittest
def set_zeros(matrix):
# don't do in place, because entire matrix
# can end up with zeros
cols = set()
rows = set()
for row_index, row in enumerate(matrix):
for col_index, val in enumerate(row):
#keep track of rows and columns with 0
if val == 0:
rows.add(row_index)
cols.add(col_index)
# traverse matrix and set coord to zero
# if row or col is in set
for row_index, row in enumerate(matrix):
for col_index, val in enumerate(row):
if row_index in rows or col_index in cols:
matrix[row_index][col_index] = 0
return matrix
class Test(unittest.TestCase):
m1 = [
[0,1,1],
[1,1,1],
[1,1,1]
]
m1_after = [
[0,0,0],
[0,1,1],
[0,1,1]
]
m2 = [
[0,1,1,1],
[1,1,1,1],
[1,1,1,0]
]
m2_after = [
[0,0,0,0],
[0,1,1,0],
[0,0,0,0]
]
m3 = [
[1,1,1,1],
[1,1,0,1],
[1,1,1,1]
]
m3_after = [
[1,1,0,1],
[0,0,0,0],
[1,1,0,1]
]
def test_set_zeros(self):
self.assertEqual(set_zeros(Test.m1),Test.m1_after)
self.assertEqual(set_zeros(Test.m2),Test.m2_after)
self.assertEqual(set_zeros(Test.m3),Test.m3_after)
if __name__ == '__main__':
unittest.main()

- We can improve the space complexity to O(1)
- Use the first row and the first col as flags
- these flags indicate whether that row or col should be zero
- Need to separately check if first row or col is zero

import unittest
def set_zeros2(matrix):
#Time Complexity: O(MN)
#Space Complexity: O(1)
if not matrix:
raise Exception("Matrix is empty")
rows = len(matrix)
cols = len(matrix[0])
first_row_zero = False
first_col_zero = False
# check if zero is in first row
if 0 in matrix[0]:
first_row_zero = True
# check if zero is in first column
for row in matrix:
if row[0] == 0:
first_col_zero = True
break
for row_index in xrange(1,rows):
for col_index in xrange(1, cols):
# use the first row and the first col
# to tell us if that row or column should
# be set to zero
if matrix[row_index][col_index] == 0:
matrix[0][col_index] = 0
matrix[row_index][0] = 0
# go through each row
for row_index in xrange(1, rows):
#check if the first item of each row is zero
if matrix[row_index][0] == 0:
#if it is, set each item in that row to 0
for col_index in xrange(1, cols):
matrix[row_index][col_index] = 0
# go through each column
for col_index in xrange(1, cols):
# check if first item in col is zero
if matrix[0][col_index] == 0:
# if it is, set each item in col to 0
for row_index in xrange(1, rows):
matrix[row_index][col_index] = 0
# if first row was flagged to be zero set
# all items in it to 0
if first_row_zero:
for col_index in xrange(0,cols):
matrix[0][col_index] = 0
# if first col was flagged to be zero set
# all items in it to 0
if first_col_zero:
for row_index in xrange(0,rows):
matrix[row_index][0] = 0
return matrix
class Test(unittest.TestCase):
m1 = [
[0,1,1],
[1,1,1],
[1,1,1]
]
m1_after = [
[0,0,0],
[0,1,1],
[0,1,1]
]
m2 = [
[0,1,1,1],
[1,1,1,1],
[1,1,1,0]
]
m2_after = [
[0,0,0,0],
[0,1,1,0],
[0,0,0,0]
]
m3 = [
[1,1,1,1],
[1,1,0,1],
[1,1,1,1]
]
m3_after = [
[1,1,0,1],
[0,0,0,0],
[1,1,0,1]
]
def test_set_zeros2(self):
self.assertEqual(set_zeros2(Test.m1),Test.m1_after)
self.assertEqual(set_zeros2(Test.m2),Test.m2_after)
self.assertEqual(set_zeros2(Test.m3),Test.m3_after)
if __name__ == '__main__':
unittest.main()