311. Sparse Matrix Multiplication
- Total Accepted: 13356
- Total Submissions: 26970
- Difficulty: Medium
- Contributors: Admin
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.
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.