Leetcode 311: Sparse Matrix Multiplication

 

Code

from collections import defaultdict

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        C = [[0] * len(B[0]) for i in xrange(len(A))]
        
        # cache for A
        dA = defaultdict(lambda: defaultdict(int))
        for i in xrange(len(A)):
            for j in xrange(len(A[0])):
                if A[i][j]:
                    dA[i][j] = A[i][j]

        # cache for B
        dB = defaultdict(lambda: defaultdict(int))
        for i in xrange(len(B)):
            for j in xrange(len(B[0])):
                if B[i][j]:
                    dB[i][j] = B[i][j]

        # only calculate necessary elements
        for i in dA:
            for j, Aij in dA[i].iteritems():
                for k, Bjk in dB[j].iteritems():
                    C[i][k] += Aij * Bjk

        return C

 

Idea

If you calculate $latex C_{ik} = \sum\limits_j A_{ij} \cdot B_{jk}$, the total time complexity is O(N^3) (assume A, B and C have size O(N^2).)

To speed it up, you need to cache non-zero of A and B. And calculate $latex A_{ij} \cdot B_{jk}$ when necessary.