LeetCode [311] 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:

Input:

A = [
  [ 1, 0, 0],
  [-1, 0, 3]
]

B = [
  [ 7, 0, 0 ],
  [ 0, 0, 0 ],
  [ 0, 0, 1 ]
]

Output:

     |  1 0 0 |   | 7 0 0 |   |  7 0 0 |
AB = | -1 0 3 | x | 0 0 0 | = | -7 0 3 |
                  | 0 0 1 |

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41

class Solution { public: vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) { unordered_map<int, unordered_set<int>> A1, B1; int ma = A.size(), na = A[0].size(); int mb = B.size(), nb = B[0].size(); for(int i=0; i<ma; ++i){ for(int j=0; j<na; ++j){ if(A[i][j]){ A1[i].insert(j); } } } for(int j=0; j<nb; ++j){ for(int i=0; i<mb; ++i){ if(B[i][j]){ B1[j].insert(i); } } } vector<vector<int>> ret(ma, vector<int>(nb, 0)); for(int i=0; i<ma; ++i){ for(int j=0; j<nb; ++j){ int s = 0; if(A1.count(i) && B1.count(j)){ for(auto k:A1[i]){ if(B1[j].count(k)){ s += A[i][k]*B[k][j]; } } } ret[i][j] = s; } } return ret; } };

//Java
class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int am = A.length, an = A[0].length;
        int bm = B.length, bn = B[0].length;
        List<Map<Integer, Integer>> mapA = new ArrayList<Map<Integer, Integer>>(am);
        List<Map<Integer, Integer>> mapB = new ArrayList<Map<Integer, Integer>>(bn);

        for(int i=0; i<am; ++i){
            mapA.add(new HashMap<>());
            for(int j=0; j<an; ++j){
                if(A[i][j]!=0){
                    mapA.get(i).put(j, A[i][j]);
                }
            }
        }

        for(int j=0; j<bn; ++j){
            mapB.add(new HashMap<>());
            for(int i=0; i<bm; ++i){
                if(B[i][j]!=0){
                    mapB.get(j).put(i, B[i][j]);
                }
            }
        }

        int[][] ret = new int[am][bn];
        for(int i=0; i<am; ++i){
            for(int j=0; j<bn; ++j){
                int s = 0;
                for(Map.Entry<Integer, Integer> e : mapA.get(i).entrySet()){
                    int k = e.getKey(), v = e.getValue();
                    if(mapB.get(j).containsKey(k)){
                        s += v*mapB.get(j).get(k);
                    }
                }
                ret[i][j] = s;
            }
        }
        return ret;
    }
}