Leetcode: Set Matrix Zeros

Given a m x n matrix, if an element is 0, set its entire row and column to 0. Do it in place.
Follow up:

Did you use extra space?
A straight forward solution using O(mn) space is probably a bad idea.
A simple improvement uses O(m + n) space, but still not the best solution.
Could you devise a constant space solution?

Understand the problem:
The problem looks quite simple at the first glance. It just needs to traverse the matrix, when we see a zero, we set the entire row and column zero. However, the problem is if we set the entire row and column zero, when we meet the zero later, we will set other rows and columns as other. And we could imagine that all matrix elements will become zero in short.
How can we get rid of this problem?

Naive Solution:
A native solution is to allocate another matrix, when we meet a zero element, we mark the entire row and column zero in the new array. However, it will take O(mn) space.

A Better Solution:
Do we really need an entire array to store the zero elements? Not really. We only need to know row i is zeros and column j is zeros, etc.. Therefore, instead of allocating an  m * n array, we only need to allocate two arrays, one with size m to store the row zero information, i.e, indicate which row is zero, the other with size n to indicate which column is zero. At last, we only need to copy those zeros into the original array. It is obviously to see that the space complexity now becomes O(m + n).

Code (Java):

public class Solution {
    public void setZeroes(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return;
        
        // get rows and columns of the matrix
        int m = matrix.length;
        int n = matrix[0].length;
        
        // allocate two addtional space to store zeros
        boolean[] row = new boolean[m];
        boolean[] col = new boolean[n];
        
        // iterate the matrix and set the zero flags
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 0) {
                    row[i] = true;
                    col[j] = true;
                }
            }
        }
        
        // set the corresponding row and columns of the original matrix as zeros
        for (int i = 0; i < m; i++) {
            if (row[i] == true) {
                for (int j = 0; j < n; j++) {
                    matrix[i][j] = 0;
                }
            }
        }
        for (int j = 0; j < n; j++) {
            if (col[j] == true) {
                for (int i = 0; i < m; i++) {
                    matrix[i][j] = 0;
                }
            }
        }
    }
}

Constant Space Solution:
Since the problem asks for a constant space solution, we must think of one without allocating additional space. Remember in last solution, we used additional row and column to store the zero information. How about we use the first column and row to save the zero information? It sounds doable. But what about the first or column contains zeros thus needs to be set as zeros as well? We could work around by using a flag to mark the first row or column as zeros. Consequently, the solution can be devised into following steps:
1. Check the first column and row contains zeros, if yes, mark the flag as true to indicate first row and column need to set to zeros.
2. Starting from matrix[1][1], check if each element  is zero, if yes, mark the corresponding first row and column elements as zeros.
3. Check back the first column and row, and set the corresponding elements in the matrix as zero.
4. Check the first row and column zero flag, if needed, set the first row and column zeros.


Code (Java):

public class Solution {
    public void setZeroes(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return;
        
        int m = matrix.length;
        int n = matrix[0].length;
        
        boolean firstRowZero = false;
        boolean firstColZero = false;
        
        // check if the first row contains zero
        for (int i = 0; i < n; i++) {
            if (matrix[0][i] == 0) {
                firstRowZero = true;
            }
        }
        
        // check if the first col contains zero
        for (int i = 0; i < m; i++) {
            if (matrix[i][0] == 0) {
                firstColZero = true;
            }
        }
        
        // check the rest of the matrix elements
        for (int i = 1; i < m; i++) {
            for (int j = 1; j < n; j++) {
                if (matrix[i][j] == 0) {
                    matrix[0][j] = 0;
                    matrix[i][0] = 0;
                }
            }
        }
        
        // set the corresponding row and col as zero
        for (int i = 1; i < m; i++) {
            for (int j = 1; j < n; j++) {
                if (matrix[i][0] == 0 || matrix[0][j] == 0) {
                    matrix[i][j] = 0;
                }
            }
        }
        
        // check and set the first row and col as zeros, if necessary
        if (firstRowZero) {
            for (int i = 0; i < n; i++) {
                matrix[0][i] = 0;
            }
        }
        
        if (firstColZero) {
            for (int j = 0; j < m; j++) {
                matrix[j][0] = 0;
            }
        }
        
    }
}

Discussion:
In this solution, since we only used two additional flag to mark the first row and column as zeros, the space complexity is constant. 

Summary:
In this post, we learned how to set zeros in a 2D matrix. Be careful about the pitfall of the problem. You cannot just set zeros in place else all matrix will become zeros eventually. We introduced three methods to reduce space cost. The third method is tricky, which used the first row and col as marker. Make sure you understand the details and try not to just memorize the solution.