311. Sparse Matrix Multiplication

https://leetcode.com/problems/sparse-matrix-multiplication/#/description

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 |

Sol:

 

 

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        
        mA = len(A)
        nA = len(A[0])
        nB = len(B[0])
        res = [[0] * len(B[0]) for _ in range(mA)]
        for row in range(mA):
            for col in range(nA):
                if A[row][col]:
                    for colB in range(nB):
                        res[row][colB] += A[row][col] * B[col][colB]
        return res