LeetCode Q311: Sparse Matrix Multiplication

Given two sparse matrices A and B, return the result of AB.

You may assume that A's column number is equal to B's row number.

Example:

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 |

Solution: Use to arrays of array to store only non-zero elements of A row-wisely and non-zero elements of B col-wisely. However, the test cases of this question in Leetcode is too weak that even the normal math solution is much faster than the answers in the below.

	class Solution {
	public:
	vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
	vector<vector<int> > res(A.size(), vector<int>(B[0].size(), 0));
	vector<vector<pair<int, int> > > rowsA(A.size());
	vector<vector<pair<int, int> > > colsB(B[0].size());
	for(int i=0; i<A.size(); i++)
	for(int j=0; j<A[0].size(); j++)
	if(A[i][j]!=0)
	rowsA[i].push_back(pair<int, int> (j, A[i][j]));
	for(int j=0; j<B[0].size(); j++)
	for(int i=0; i<B.size(); i++)
	if(B[i][j]!=0)
	colsB[j].push_back(pair<int, int> (i, B[i][j]));
	for(int i=0; i<A.size(); i++){
	for(int j=0; j<B[0].size(); j++){
	int p=0, q=0;
	int sum =0 ;
	while(p!=rowsA[i].size()&&q!=colsB[j].size()){
	if(rowsA[i][p].first==colsB[j][q].first){
	sum = sum + rowsA[i][p].second*colsB[j][q].second;
	p++;
	q++;
	continue;
	}
	if(rowsA[i][p].first<colsB[j][q].first)
	p++;
	if(rowsA[i][p].first>colsB[j][q].first)
	q++;
	}
	res[i][j]=sum;
	}
	}
	return res;
	}
	};

view raw Q311.cpp hosted with ❤ by GitHub