Sparse Matrix Multiplication

sparse_matrix.py

"""
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:
Input:
A = [
  [ 1, 0, 0],
  [-1, 0, 3]
]
B = [
  [ 7, 0, 0 ],
  [ 0, 0, 0 ],
  [ 0, 0, 1 ]
]
Output:
AB = [
  [7, 0, 0],
  [-7, 0, 3]
]
"""

class SparseMatrix:
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        if not A or not A[0] or not B or not B[0]: return None
        m, n, l = len(A), len(A[0]), len(B[0])
        result = [[0] * l for _ in range(m)]
        a, b = [{} for _ in range(m)], [{} for _ in range(l)]
        # [{col: val}, ...]
        for i in range(m):
            for j in range(n):
                if A[i][j] != 0:
                    a[i][j] = A[i][j]
        for i in range(n):
            for j in range(l):
                if B[i][j] != 0:
                    b[j][i] = B[i][j]
        for i in range(m):
            for j in range(l):
                for k in a[i]:
                    if k in b[j]:
                        result[i][j] += a[i][k] * b[j][k]
        return result