2
\$\begingroup\$

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()
\$\endgroup\$
2
  • \$\begingroup\$ Is numpy not allowed? \$\endgroup\$ Sep 19, 2016 at 18:34
  • \$\begingroup\$ The problem does not say you can not use numpy. I am solving these problems to learn programming and to prepare for technical interviews. \$\endgroup\$
    – Average
    Sep 19, 2016 at 19:11

2 Answers 2

1
\$\begingroup\$

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

\$\endgroup\$
4
  • \$\begingroup\$ Why do we need two for rowi loops? \$\endgroup\$
    – Average
    Sep 19, 2016 at 19:20
  • \$\begingroup\$ @newToProgramming: typo, sorry, fixed \$\endgroup\$ Sep 19, 2016 at 19:36
  • \$\begingroup\$ ~all(row) isn't a good idea. ~ is a bitwise not, not a logical not, and since True == 1, ~True == -2, meaning that bool(True) == bool(~True) == True. \$\endgroup\$
    – DSM
    Sep 3, 2018 at 14:52
  • \$\begingroup\$ @DSM Good catch, fixed \$\endgroup\$ Sep 6, 2018 at 15:56
1
\$\begingroup\$

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.

\$\endgroup\$
3
  • \$\begingroup\$ thank you. I will try and understand the space and time complexity of your implementation of locate_zeros, I assume it is O(1) space O(m*n) time? Though in retrospect i should have combined the two functions into one even without implementing it as a generator I assume. \$\endgroup\$
    – Average
    Sep 18, 2016 at 22:47
  • \$\begingroup\$ What is zeros_positions? \$\endgroup\$
    – vnp
    Sep 18, 2016 at 22:55
  • \$\begingroup\$ @vnp locate_zeros in the matrix \$\endgroup\$
    – Caridorc
    Sep 19, 2016 at 12:34

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.