654. Sparse Matrix Multiplication

Code(Language:C++)

class Solution {
public:
    /**
     * @param A: a sparse matrix
     * @param B: a sparse matrix
     * @return: the result of A * B
     */
    vector<vector<int>> multiply(vector<vector<int>> &A, vector<vector<int>> &B) {
        // write your code here
        vector<vector<int>> res;
        const int mA = A.size();
        const int mB = B.size(); 
        if(mA == 0 || mB == 0){
            return res; 
        }
        const int nA = A[0].size(); 
        const int nB = B[0].size(); 
        if(nA != mB){
            return res; 
        }
        // search non 0 in each row of B;
        vector<vector<int>> non0B(mB, vector<int>());
        for(int j = 0; j < mB; j++){
            for(int i = 0; i < nB; i++){
                if(B[j][i]){
                    non0B[j].push_back(i); 
                }
            }
        }
        
        res = vector<vector<int>>(mA, vector<int>(nB, 0)); 
        for(int i = 0; i < mA; i++){
            for(int k = 0; k < nA; k++){
                if(A[i][k] == 0){
                    continue; 
                }
                for(auto j : non0B[k]){
                    res[i][j] += A[i][k] * B[k][j]; 
                }
                
            }
        }
        return res; 
    }
};