interview-questions/arrays-and-strings/matrix_zeros.py

# Write an algorithm such that if an element in an MxN matrix is 0,
# its entire row and column is set to 0.

import pytest
import numpy as np
from numpy.testing import assert_array_equal

def test_matrix_zeros():
    case1 = np.zeros((3,3))
    case2 = np.array([[1,0,1],[2,3,0]])
    expect2 = np.zeros((2,3))
    assert_array_equal(matrix_zeros(case1), case1)
    assert_array_equal(matrix_zeros(case2), expect2)
    # test error message
    with pytest.raises(Exception):
        matrix_zeros([1,1])

def matrix_zeros(matrix):
    r, c = matrix.shape
    if not r or not c:
        raise ValueError('Not a matrix!')

    r_ind, c_ind = [], []
    for i in range(r):
        for j in range(c):
            if matrix[i,j] == 0:
                r_ind.append(i)
                c_ind.append(j)
    for i in r_ind:
        matrix[i, :] = np.zeros(c)
    for i in c_ind:
        matrix[:, i] = np.zeros(r)
    return matrix