【LeetCode】311. 稀疏矩阵的乘法

本题是 LeetCode 会员才能看…

一、题目描述

给你两个 稀疏矩阵 A 和 B，请你返回 AB 的结果。 你可以默认 A 的列数等于 B 的行数。

请仔细阅读下面的示例。

示例： 输入： A = [ [ 1, 0, 0], [-1, 0, 3] ] B = [ [ 7, 0, 0 ], [ 0, 0, 0 ], [ 0, 0, 1 ] ] 输出： | 1 0 0 | | 7 0 0 | | 7 0 0 | AB = | -1 0 3 | x | 0 0 0 | = | -7 0 3 | | 0 0 1 |

二、解题思路 & 代码

2.1 暴力求解

没有根据稀疏矩阵的特性进行优化

def matrixMul1(A, B): """暴力法""" m = len(A) n = len(A[0]) p = len(B[0]) ans = [[0] * p for i in range(m)] for i in range(m): for j in range(p): sum_ = 0 for k in range(n): sum_ += A[i][k] * B[k][j] ans[i][j] = sum_ return ans

2.2 优化（选取都不为0的行和列相乘）

稀疏矩阵就是有很多 0 ，所以其实优化就是要把都是 0 的行和列去掉

def matrixMul2(A, B): m = len(A) n = len(A[0]) p = len(B[0]) r_all0 = [True] * m # 标记全为 0 的 row c_all0 = [True] * p # 标记全为 0 的 col # 用 flag 做标记 flag = False for i in range(m): flag = False for j in range(n): if A[i][j]: flag = True break if flag: r_all0[i] = False for j in range(p): flag = False for i in range(n): if B[i][j]: flag = True break if flag: c_all0[j] = False ans = [[0] * p for i in range(m)] for i in range(m): for j in range(p): if (r_all0[i] or c_all0[j]): # 优化在这里 continue sum_ = 0 for k in range(n): sum_ += A[i][k] * B[k][j] ans[i][j] = sum_ return ans

参考：

LeetCode 311. 稀疏矩阵的乘法