[leetcode] 311. Sparse Matrix Multiplication 解题报告

题目链接: 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;
    }
};