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?

---

To achieve constant space, we need to store the information of which row and which column will be set to somewhere already exists. Following this, we can use the first row to store which column need to be set and first column to store which row need to be set.

C++ Code:

/*
 * func: set_matrix_zero
 * goal: if an element is 0, set its entire row and column to 0
 * @param matrix: input matrix
 * return:
 */
/*
 * check every position of 0 and store the info into first row and
 * first column of the matrix.
 * complexity: time O(mn), space O(1)
 */
void set_matrix_zero(vector<vector<int> > &matrix){
    int m = matrix.size();
    if(m == 0){
        return;
    }
    int n = matrix[0].size();
    bool first_row = false;
    bool first_column = false;
    //Check if there is 0 in first row or first column
    for(int i = 0; i < m; ++i){
        if(matrix[i][0] == 0){
            first_column = true;
            break;
        }
    }
    for(int j = 0; j < n; ++j){
        if(matrix[0][j] == 0){
            first_row = true;
            break;
        }
    }
    
    //Find all 0s in the matrix
    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 row 0s
    for(int i = 1; i < m; ++i){
        if(matrix[i][0] == 0){
            for(int j = 1; j < n; ++j){
                matrix[i][j] = 0;
            }
        }
    }
    
    //Set column 0s
    for(int j = 1; j < n; ++j){
        if(matrix[0][j] == 0){
            for(int i = 1; i < m; ++i){
                matrix[i][j] = 0;
            }
        }
    }
    
    //Set first row if needed
    if(first_row){
        for(int j = 0; j < n; ++j)
            matrix[0][j] = 0;
    }
    
    //Set first column if needed
    if(first_column){
        for(int i = 0; i < m; ++i)
            matrix[i][0] = 0;
    }
}

Python Code:

# func: if an element is 0, set its entire row and column to 0
# @param matrix: input matrix
# @return:
# complexity: time O(mn), space O(1)
def set_matrix_zero(matrix):
    m = len(matrix)
    if m == 0:
        return

    n = len(matrix[0])
    first_row = False
    first_column = False

    #Check if first row need to be set
    for j in xrange(n):
        if matrix[0][j] == 0:
            first_row = True
            break
    #Check if first column need to be set
    for i in xrange(m):
        if matrix[i][0] == 0:
            first_column = True
            break

    #Check remaining elements
    for i in xrange(1, m):
        for j in xrange(1, n):
            if matrix[i][j] == 0:
                matrix[i][0] = 0
                matrix[0][j] = 0

    #Set elements
    for i in xrange(1, m):
        for j in xrange(1, n):
            if matrix[i][0] == 0 or matrix[0][j] == 0:
                matrix[i][j] = 0

    #Set first row
    if first_row:
        matrix[0] = [0] * n

    #Set first column
    if first_column:
        for j in xrange(m):
            matrix[j][0] = 0