题目
题干
该问题稀疏矩阵乘法 题面:
Given two sparse matrices A and B, return the result of AB.
给你两个 稀疏矩阵 A 和 B,请你返回 AB 的结果。
你可以默认 A 的列数等于 B 的行数。
示例
输入:
A = [[1,0,0],[-1,0,3]]
B = [[7,0,0],[0,0,0],[0,0,1]]
输出:
[[7,0,0], [-7,0,3]]
题解
由于是稀疏数组可以只计算非零数字。
将A矩阵转存为hashmap,然后对每一个非零数遍历B的列数来更新它在ans矩阵里可以被用到的位置。
C++
class Solution {
public:
vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
int n = A.size(), m = A[0].size(), k = B[0].size();
vector<vector<int>> res(n, vector<int>(k, 0));
unordered_map<int, int> am;
//存储A的非零元素
for (int i = 0; i < n; i ++)
for (int j = 0; j < m; j ++)
{
if (!A[i][j]) continue;
int idx = i * m + j;
am[idx] = A[i][j];
}
for (auto aij : am)
{
int va = aij.second, aidx = aij.first, ai = aidx / m, aj = aidx % m;
for (int j = 0; j < k; j ++)
{
res[ai][j] += va * B[aj][j];
}
}
return res;
}
};
Python
class Solution:
def __init__(self,A,B) -> None:
self.a = A
self.b = B
def multiply(self):
am = {
}
n = len(self.a)
m = len(self.a[0])
k = len(self.b[0])
res = []
#存储A的非零元素
for i in range(n):
res = res+[[0]*k]
for i in range(n):
for j in range(m):
if self.a[i][j] == 0:
continue
idx = i * m + j
am[idx] = self.a[i][j]
for aidx, va in am.items():
ai = aidx // m
aj = aidx % m
for j in range(k):
res[ai][j] += va * self.b[aj][j]
return res