稀疏矩阵的乘法

最近做一个题， 稀疏矩阵的乘法， 自己想了下，最直观的想法当然是 记录 下来 哪些 不为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;
}