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:

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 |

Solution 1: Only do the calculation for non-zero elements in A

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        rowA = len(A)
        colA = len(A[0])
        colB = len(B[0])
        ret = [[0 for column in range(colB)] for row in range(rowA)]
        for i in range(rowA):
            for j in range(colA):
                if A[i][j] != 0:
                    for k in range(colB):
                        ret[i][k] += A[i][j] * B[j][k]
        return ret

Solution 2. Using dictionary to store non-zero elements in A

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        m = len(A)
        l = len(A[0])
        n = len(B[0])
        ret = [[0 for col in range(n)] for row in range(m)]
        dic = {}
        # final all nonzero elements in A
        for i in range(m):
            tmp = {}
            for j in range(l):
                if A[i][j] != 0:
                    tmp[j] = A[i][j]
            dic[i] = tmp
        # compute matrix multiplication
        for i in range(m):
            tmp = dic.get(i)
            for j, val in tmp.items():
                for k in range(n):
                    ret[i][k] += val * B[j][k]
        return ret