Set Matrix Zeroes

Problem

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?

Examples

Example 1:

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

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

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

You should aim for a solution with O(m*n) time and O(1) space, where m is the number of rows and n is the number of columns in the given matrix.

Solution

We can’t solve this in a single pass, as we don’t know what zeros existed before versus ones we just set. We can store the rows/cols that need to be zeroed, then do another pass. To do this in O(1) space, the first row/col holds if the corresponding col/row needs to be zeroed. We need 2 variables, one for first_row and first_col to know if we need to zero out the first row/col. Then we do a second pass, zeroing out the elements.

class Solution {
public:
    void setZeroes(vector<vector<int>>& matrix) {
        const int M = matrix.size();
        const int N = matrix[0].size();
        bool first_row = false;
        bool first_col = false;
 
        // Scan for zeros, and store in first row/col
        // If zero occurs in first row/col, we store seprately
        for (int r = 0; r < M; ++r) {
            for (int c = 0; c < N; ++c) {
                if (matrix[r][c] == 0) {
                    if (r > 0) { matrix[r][0] = 0; }
                    else { first_row = true; }
                    if (c > 0) { matrix[0][c] = 0; }
                    else { first_col = true; }
                }
            }
        }
 
        // Scan top row and column and set corresponding col/row 
        for (int r = 1; r < M; ++r) {
            if (matrix[r][0] == 0) {
                for (int c = 1; c < N; ++c) {
                    matrix[r][c] = 0;
                }
            }
        }
        for (int c = 1; c < N; ++c) {
            if (matrix[0][c] == 0) {
                for (int r = 1; r < M; ++r) {
                    matrix[r][c] = 0;
                }
            }
        }
 
        // If first row is zero'd do it
        if (first_row) {
            for (int c = 0; c < N; ++c) {
                matrix[0][c] = 0;
            }
        }
        // If first col is zero'd do it
        if (first_col) {
            for (int r = 0; r < M; ++r) {
                matrix[r][0] = 0;
            }
        }
    }
};