【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      

2.2 優化（選取都不為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. 稀疏矩陣的乘法​​