Problem statement

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

Solution

First, for every row of first matrix and for every column of second matrix, we need to create hash-table: with keys are indexes for non-zero values and values are corresponding elements for these elements. Then when we multiply row and column, we need to iterate over two hash-tables and find sum faster than O(n).

Complexity

Overall time complexity is < O(m*n*k), because we have sparse matrixes. However it is difficult to say like this, what is exact complexity

Code

class Solution:
    def multiply(self, mat1, mat2):
        def process_matrix(mat):
            d = defaultdict(dict)
            for i, row in enumerate(mat):
                for j, elem in enumerate(row):
                    if elem != 0: d[i][j] = elem
            return d
        
        m1 = process_matrix(mat1)
        m2 = process_matrix(zip(*mat2))
        
        m, n, k = len(mat1), len(mat2[0]), len(mat2)
        
        ans = [[0] * n for _ in range(m)]
        
        for i in range(m):
            for j in range(k):
                for s in m1[i]:
                    if s in m2[j]:
                        ans[i][j] += m1[i][s]*m2[j][s]
        
        return ans