本题是 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

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

参考：

  1. ​​LeetCode 311. 稀疏矩阵的乘法​​