LeetCode 311. Sparse Matrix Multiplication

需要利用 sparse 的特性来做这道题。扫描A的每一个元素 A[i][j],考虑A[i][j]会被B中哪些元素乘到。

对于A的第i行 A[i][*],会乘以B的每一列 B[*][k] ∀k,得到 res[i][k]。因此,最后 res[i][k] += A[i][j] * B[j][k]。

这样的好处是,可以先判断 A[i][j] 是否为0,如果是0,可以节省所有需要和 A[i][j] 相乘的运算。

class Solution {
public:
    vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
        int m1=A.size(), n1=A[0].size(), m2=B.size(), n2=B[0].size();
        vector<vector<int>> res;
        if (m1==0 || n1==0 || m2==0 || n2==0 || n1!=m2) return res;
        
        res.resize(m1,vector<int>(n2,0));
        for (int i=0;i<m1;++i){
            for (int j=0;j<n1;++j){
                if (A[i][j]){
                    for (int k=0;k<n2;++k){
                        if (B[j][k]){
                            res[i][k] += A[i][j] * B[j][k];
                        }
                    }
                }
            }
        }
        return res;
    }
};

猜你喜欢

转载自www.cnblogs.com/hankunyan/p/11257614.html