程序员面试金典 - 面试题 01.08. 零矩阵

题目难度: 中等

原题链接

今天继续更新程序员面试金典系列, 大家在公众号 算法精选 里回复 面试金典 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

编写一种算法，若 M × N 矩阵中某个元素为 0，则将其所在的行与列清零。

示例 1：

输入：
[
  [1,1,1],
  [1,0,1],
  [1,1,1]
]
输出：
[
  [1,0,1],
  [0,0,0],
  [1,0,1]
]

示例 2：

输入：
[
  [0,1,2,0],
  [3,4,5,2],
  [1,3,1,5]
]
输出：
[
  [0,0,0,0],
  [0,4,5,0],
  [0,3,1,0]
]

题目思考

如何做到不使用额外空间?

解决方案

方案 1

思路

分析题目, 直观的思路是两次遍历
第一次遍历, 统计所有有 0 的行和列分别放入两个集合中
第二次遍历, 根据当前行和列是否在集合中来决定是否要将整行或者整列清 0

复杂度

时间复杂度 O(M*N): 需要遍历矩阵两遍
空间复杂度 O(M+N): 需要使用行列两个集合

代码

class Solution:
    def setZeroes(self, matrix: List[List[int]]) -> None:
        """
        Do not return anything, modify matrix in-place instead.
        """
        # 方法1: O(M*N)时间, O(M+N)空间
        # 使用两个集合存储有0的行或列
        if not matrix:
            return
        rows, cols = len(matrix), len(matrix[0])
        zeroRows, zeroCols = set(), set()
        for r in range(rows):
            for c in range(cols):
                if matrix[r][c] == 0:
                    zeroRows.add(r)
                    zeroCols.add(c)
        for r in zeroRows:
            matrix[r] = [0] * cols
        for c in zeroCols:
            for r in range(rows):
                matrix[r][c] = 0

方案 2

思路

方案 1 效率已经很高了, 但有没有办法不使用额外空间来存储行列为零的标记呢?
答案也是有的, 我们可以利用原始矩阵, 直接将某行或某列是否为 0 的信息存储在第 0 行或第 0 列即可
不过这样做我们需要特殊处理第 0 行/第 0 列本身有元素为 0 的情况, 可以使用两个 bool 来先判断是否存在这样的情况, 之后再进行其他行列的 0 标记
下面代码中对每一步骤有更加详细的注释, 方便大家理解

复杂度

时间复杂度 O(M*N): 需要遍历矩阵两遍
空间复杂度 O(1): 只使用了几个变量

代码

class Solution:
    def setZeroes(self, matrix: List[List[int]]) -> None:
        """
        Do not return anything, modify matrix in-place instead.
        """
        # 方法2: O(M*N)时间, O(1)空间
        # 如果某个位置为0, 则将其同一行列的第0个元素标成0
        # 最后再统一把标0的行和列全变成0
        # 注意边界, 不能变第0行/列以免污染标记
        # 需要提前判断第0行/列是否应该全成0, 是的话在最后再标记
        if not matrix:
            return
        rows, cols = len(matrix), len(matrix[0])
        # 使用两个bool, 特殊处理第0行/第0列本身有元素为0的情况
        startColZero = any([matrix[r][0] == 0 for r in range(rows)])
        startRowZero = any([matrix[0][c] == 0 for c in range(cols)])
        # 将行/列的0信息存放在对应的第0行/第0列中
        for r in range(rows):
            for c in range(cols):
                if matrix[r][c] == 0:
                    matrix[r][0] = matrix[0][c] = 0
        # 标记除了第0行和第0列的其他部分
        for r in range(1, rows):
            if matrix[r][0] == 0:
                matrix[r] = [0] * cols
        for c in range(1, cols):
            if matrix[0][c] == 0:
                for r in range(1, rows):
                    matrix[r][c] = 0
        # 最后根据上面得到的两个bool标记第0行和第0列, 注意这一步不能在上一步之前进行, 不然就污染了之前存储好的在第0行或列的标记了
        if startColZero:
            for r in range(rows):
                matrix[r][0] = 0
        if startRowZero:
            matrix[0] = [0] * cols