311. Sparse Matrix Multiplication - Medium

Posted 2021-01-24 fatttcat

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 |

 

M1: brute force + 忽略0

时间：O(N^3)，空间：O(1)

class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int m = A.length, n = A[0].length, l = B[0].length;  // B.length == n
        int[][] AB = new int[m][l];
        for(int i = 0; i < m; i++) {
            for(int j = 0; j < n; j++) {
                if(A[i][j] == 0) continue;
                for(int k = 0; k < l; k++) {
                    if(B[j][k] != 0)
                        AB[i][k] += A[i][j] * B[j][k];
                }
            }
        }
        return AB;
    }
}

 

M2: hash table

用两个hashmap分别存A，B中的非零元素，记下非零元素及其行、列号。然后再遍历A的keyset，如果在mapB中存在与之对应相乘的元素，遍历B得到乘加和。因为忽略了大量0，当矩阵很稀疏时，将非0元素存入hashmap的时间相对乘加运算可以忽略，因此速度可以得到提高。

时间：O(N^3)，空间：O(N)

class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        HashMap<Integer, HashMap<Integer, Integer>> mapA = new HashMap<>();  // <i, <j, A[i][j]>>
        HashMap<Integer, HashMap<Integer, Integer>> mapB = new HashMap<>();
        
        for(int i = 0; i < A.length; i++) {
            for(int j = 0; j < A[0].length; j++) {
                if(A[i][j] != 0) {
                    mapA.putIfAbsent(i, new HashMap<Integer, Integer>());
                    mapA.get(i).put(j, A[i][j]);
                }
            }
        }
        for(int i = 0; i < B.length; i++) {
            for(int j = 0; j < B[0].length; j++) {
                if(B[i][j] != 0) {
                    mapB.putIfAbsent(i, new HashMap<Integer, Integer>());
                    mapB.get(i).put(j, B[i][j]);
                }
            }
        }
        
        int[][] AB = new int[A.length][B[0].length];
        for(int i : mapA.keySet()) {
            for(int j : mapA.get(i).keySet()) {
                if(mapB.containsKey(j)) {
                    for(int k : mapB.get(j).keySet()) {
                        AB[i][k] += mapA.get(i).get(j) * mapB.get(j).get(k);
                    }
                }
            }
        }
        return AB;
    }
}