算法题--将矩阵中0所在的行和列置0

image.png

0. 链接

题目链接

1. 题目

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

Example 1:

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

Example 2:

Input: 
[
  [0,1,2,0],
  [3,4,5,2],
  [1,3,1,5]
]
Output: 
[
  [0,0,0,0],
  [0,4,5,0],
  [0,3,1,0]
]

Follow up:

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?

2. 思路1：两次遍历+第0行/列单独处理

1. 标志位分配

• 用first_col_zero(True/False)来记录第0列是否需要全部置0
• 用matrix[i][0]记录第i行是否需要全部置为0, 其中0 <= i <= m - 1
• 用matrix[0][j]记录第j列是否需要全部置为0, 其中1 <= j <= n - 1

2. 算法过程: 2次遍历

• 第1次通过遍历 0 <= i <= m - 1 和 1 <= j <= n - 1来设置标志位
• 第2次遍历先试图对行(1-m)且列(1-n)的格子置0，再对第0行/列则单独尝试置0

3. 代码

class Solution:
    def setZeroes(self, matrix: List[List[int]]) -> None:
        first_col_zero = False
        m = len(matrix)
        n = len(matrix[0])

        for i in range(m):
            if matrix[i][0] == 0:
                first_col_zero = True
            for j in range(1, n):
                if matrix[i][j] == 0:
                    matrix[0][j] = 0
                    matrix[i][0] = 0

        for i in range(1, m):
            for j in range(1, n):
                if matrix[i][0] == 0 or matrix[0][j] == 0:
                    matrix[i][j] = 0

        if matrix[0][0] == 0:
            for j in range(1, n):
                matrix[0][j] = 0

        if first_col_zero:
            for i in range(1, m):
                matrix[i][0] = 0


def print_matrix(matrix):
    for each in matrix:
        print(each)
    print('=' * 50)


def my_test(solution, matrix):
    print('input:')
    print_matrix(matrix)
    solution.setZeroes(matrix)
    print('output:')
    print_matrix(matrix)


solution = Solution()
my_test(solution, [
    [1, 1, 1],
    [1, 0, 1],
    [1, 1, 1]
])
my_test(solution, [
    [0, 1, 2, 0],
    [3, 4, 5, 2],
    [1, 3, 1, 5]
])

输出结果

input:
[1, 1, 1]
[1, 0, 1]
[1, 1, 1]
==================================================
output:
[1, 0, 1]
[0, 0, 0]
[1, 0, 1]
==================================================
input:
[0, 1, 2, 0]
[3, 4, 5, 2]
[1, 3, 1, 5]
==================================================
output:
[0, 0, 0, 0]
[0, 4, 5, 0]
[0, 3, 1, 0]
==================================================

4. 结果

image.png