题目链接: https://leetcode.com/problems/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 |
思路: 矩阵相差的题目, 如果直接相乘的话碰到很多0的情况会超时, 因此可以做一个简单的优化, 如果A的某一行全是0的话,则可以不用再计算了.
代码如下:
class Solution {
public:
vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
if(A.size()==0 || B.size() ==0) return {};
int m = A.size(), n = A[0].size(), p = B[0].size();
vector<vector<int>> result(m, vector<int>(p, 0));
set<int> row, col;
for(int i = 0, j; i < m; i++)
{
for(j =0; j < n; j++) if(A[i][j]) break;
if(j == n) row.insert(i);
}
for(int i = 0, j; i < p; i++)
{
for(j = 0; j < n; j++) if(B[j][i]) break;
if(j == n) col.insert(i);
}
for(int i = 0; i < m; i++)
{
if(row.count(i)) continue;
for(int j = 0; j < p; j++)
{
if(col.count(j)) continue;
for(int k = 0, sum=0; k < n; k++)
result[i][j] += A[i][k]*B[k][j];
}
}
return result;
}
};
以上是常规的一种方法,还有一种更简洁的方法.只需要改变一下三重循环的第二三重即可.
代码如下:
class Solution {
public:
vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
int m = A.size(), n = A[0].size(), p = B[0].size();
vector<vector<int>> ans(m, vector<int>(p, 0));
for(int i = 0; i < m; i++)
{
for(int k = 0; k < n; k++)
{
if(A[i][k] == 0) continue;
for(int j = 0; j < p; j++) ans[i][j] += A[i][k] * B[k][j];
}
}
return ans;
}
};