311.Sparse Matrix Multiplication

Tags: [math], [matrix], [skip]

Com: {fb}

Link: 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 |

Solution: Math, Matrix, Skip

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        num_of_a_rows = len(A)
        num_of_a_cols = len(A[0])
        num_of_b_cols = len(B[0])

        result = [[0 for _ in xrange(num_of_b_cols)] for _ in xrange(num_of_a_rows)]
        for a_row in xrange(num_of_a_rows):
            for a_col in xrange(num_of_a_cols):
                if not A[a_row][a_col]:
                    continue

                for b_col in xrange(num_of_b_cols):
                    if not B[a_col][b_col]:
                        continue

                    result[a_row][b_col] += A[a_row][a_col] * B[a_col][b_col]

        return result

Revelation:

• Usually, we do the matrix multiplication, there are three levels of loops, the first level is for a_row in xrange(num_of_a_rows), and the second level is for b_col in xrange(num_of_b_cols), then third level is for k in xrange(numof_a_cols), but the order of the loops can be changed, because we just need to fill result[a_row][b_col] one by one. So if the second level we use for a_col in xrange(num_of_a_cols), we can get A[a_row][a_col] before the third level, if A[a_row][b_col] is 0, we can just skip the whole third level loop, which will make the performance better.

Note:

• Time complexity = O(num_of_a_rows * num_of_a_cols * num_of_b_cols).

Time Limited Exceeded:

class Solution(object):
    def multiply(self, A, B):
        """
        :type A: List[List[int]]
        :type B: List[List[int]]
        :rtype: List[List[int]]
        """
        num_of_a_rows = len(A)
        num_of_a_cols = len(A[0])
        num_of_b_cols = len(B[0])

        result = [[0 for _ in xrange(num_of_b_cols)] for _ in xrange(num_of_a_rows)]
        for row in xrange(num_of_a_rows):
            for col in xrange(num_of_b_cols):
                tmp_sum = 0
                for k in xrange(num_of_a_cols):
                    tmp_sum += A[row][k] * B[k][col]

                result[row][col] = tmp_sum

        return result

Note:

• Time complexity = O(num_of_a_rows * num_of_a_cols * num_of_b_cols).