Write an algorithm such that if an element in an MxN matrix is zero, its entire row and column are set to zero.
Hint
- If you just cleared the rows and columns as you found zeroes, you'd likely wind up clearing the whole matrix. Try finding the cells with zeros first before making any changes to the matrix.
- Can you use O(N) additional space instead of O(N^2)? What information do you really need from the list of cells that are zero?
- You probably need some data storage to maintain a list of the rows and columns that need to be zeroed. Can you reduce the additional space usage to O(1) by using the matrix itself for data storage?
Solution
首先很容易想到空间复杂度为O(MN)的解,直接新建一个和原矩阵大小一样的矩阵,然后逐行遍历原矩阵,遍历过程中遇到0,就将新矩阵对应行全部设为0,接着再逐列遍历原矩阵做同样的操作。但那样复杂度太高,优化到O(M+N)的方法是用一个长度为M的数组记录各行中是否有0,再用一个长度为N的数组记录各列中是否有0,最后直接更新原矩阵。
既然题目中提到了有空间复杂度O(1)的解法,那我们就不能创建额外的数组,于是考虑使用原矩阵的第一行与第一列来做记录,过程如下:
- 先遍历第一行与第一列,如果有0,则将各自的flag设置为true
- 接着遍历除第一行与第一列外的整个矩阵,如果有0,则将对应的第一行与第一列的值设为0
- 再次遍历除第一行与第一列外的整个矩阵,如果对应的第一行与第一列的数字有一个为0,则将当前值设为0
- 最后根据第一行与第一列的flag来更新第一行与第一列
public void zeroMatrix(int[][] matrix) {
if (matrix == null || matrix.length == 0) return;
int rows = matrix.length, cols = matrix[0].length;
boolean rowZero = false, colZero = false;
for (int i = 0; i < rows; i++) {
if (matrix[i][0] == 0) colZero = true;
}
for (int i = 0; i < cols; i++) {
if (matrix[0][i] == 0) rowZero = true;
}
for (int i = 1; i < rows; i++) {
for (int j = 1; j < cols; j++) {
if (matrix[i][j] == 0) {
matrix[0][j] = 0;
matrix[i][0] = 0;
}
}
}
for (int i = 1; i < rows; i++) {
for (int j = 1; j < cols; j++) {
if (matrix[0][j] == 0 || matrix[i][0] == 0) matrix[i][j] = 0;
}
}
if (rowZero) {
for (int i = 0; i < cols; i++) matrix[0][i] = 0;
}
if (colZero) {
for (int i = 0; i < rows; i++) matrix[i][0] = 0;
}
}