[NeetCode 150] Set Zeroes in Matrix

Set Zeroes in Matrix

Given an m x n matrix of integers matrix, if an element is 0, set its entire row and column to 0’s.

You must update the matrix in-place.

Follow up: Could you solve it using O(1) space?

Example 1:

Input: matrix = [
  [0,1],
  [1,1]
]

Output: [
  [0,0],
  [0,1]
]

Example 2:

Input: matrix = [
  [1,2,3],
  [4,0,5],
  [6,7,8]
]

Output: [
  [1,0,3],
  [0,0,0],
  [6,0,8]
]

Constraints:

1 <= matrix.length, matrix[0].length <= 100
-2^31 <= matrix[i][j] <= (2^31) - 1

Solution

To achieve $O(1)$ extra space, we have to reuse the space in matrix to mark the columns we need to set to zero. So, we can keep the first row as the mark row. An additional variable is applied to mark whether the first row need to be set to zero. Then, go through other rows, if zero exists, set the corresponding column in the first row to zero, set the whole row to zero afterwards. Finally, set columns to zero according to the mark in first row.

Code

class Solution:
    def setZeroes(self, matrix: List[List[int]]) -> None:
        zero_in_first_row = False
        for i in range(len(matrix[0])):
            if matrix[0][i] == 0:
                zero_in_first_row = True
                break

        for i in range(1, len(matrix)):
            zero = False
            for j in range(len(matrix[0])):
                if matrix[i][j] == 0:
                    zero = True
                    matrix[0][j] = 0
            if zero:
                for j in range(len(matrix[0])):
                    matrix[i][j] = 0
        
        for i in range(len(matrix[0])):
            if matrix[0][i] == 0:
                for j in range(len(matrix)):
                    matrix[j][i] = 0
        
        if zero_in_first_row:
            for i in range(len(matrix[0])):
                matrix[0][i] = 0