最近做一个题, 稀疏矩阵的乘法, 自己想了下,最直观的想法当然是 记录 下来 哪些 不为0的元素,只对这些元素进行操作, 假设 A * B, 由于 矩阵乘法的定义是 A的某行乘以 B的某列, 产生的第一个想法是 记录下来 A的每一行中不为0的元素的列号, 记为 集合 AN, 记录下来 B的每一列的行号,记为 集合 BN。然后 进行计算。 等写代码的时候才发现, 需要列号 和行号对应 两个数字 才能相乘。 如果这样记录,为了避免再一次判断某个元素是否为0, 还得考虑 AN 和 BN的交集, 这样一样就比较麻烦。
想起之前看过 非线性优化库g2o的源码,里面引入了稀疏矩阵运算, 使用了两种稀疏矩阵的存储结构,思路都是是 存储每一列不为0的元素的行号,具体用的数据结构不同而已,一种是vector, 一种是map。 进行乘法的时候并没有这么麻烦, 于是再次去看了下 代码, 原来 不是一次性先把 乘积中某个元素的值求出来, 而是 遍历 AN和 BN , 算出 乘积中某个元素的一部分, 逐步累加出 乘积中某个元素的值。 这样就避免了前面求交集的麻烦!
全部代码如下,测试通过, 速度自然是刚刚的。
// 每列不为0的元素的行号
void getAllNonZero(vector<vector<int>>& A, int m, int n, vector<vector<int>> & NonZero) {
NonZero.resize(n);
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
if (A[i][j] != 0) {
NonZero[j].push_back(i);
}
}
}
}
// 稀疏矩阵乘法
vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
vector<vector<int>> _res;
int m = A.size();
if (m < 1) {
return _res;
}
int n = A[0].size();
if (n < 1) {
return _res;
}
if (n != B.size()) {
return _res;
}
int k = B[0].size();
// m*n * n*k == m*k
vector<vector<int>> res(m, vector<int>(k,0));
// 先检查A 中每列不为0的元素
// 检查 B 中 每列不为0的元素
vector<vector<int>> A_NonZero;
getAllNonZero(A, m, n, A_NonZero);
vector<vector<int>> B_NonZero;
getAllNonZero(B, n, k, B_NonZero);
// 遍历B的每一列 不为0的元素
for (int i = 0; i < k; i++) {
vector<int> & currCol = B_NonZero[i];
for (auto & Brow : currCol) {
vector<int> & ACol = A_NonZero[Brow];
for (auto & Arow : ACol) {
res[Arow][i] += A[Arow][Brow] * B[Brow][i];
}
}
}
return res;
}
int main() {
vector<vector<int>>A = {
{ 1, 0, 0 },
//{ -1, 0, 3 }
};
vector<vector<int>> B = {
{ 7, 0, 0 },
{ 0, 0, 0 },
{ 0, 0, 1 }
};
vector<vector<int>> C = multiply(A, B);
output(C);
return 0;
}