Katkat

Multiply Two Sparse Matrices

AppleSoftware Engineer · Technical Screen

Last updated

This problem tests your ability to optimize matrix multiplication by leveraging data sparsity to avoid unnecessary computations and handle memory efficiently.

Multiply Two Sparse Matrices

Given two sparse matrices A and B, represented as 2D arrays, compute their product C = A * B. A matrix is sparse if most of its elements are zero. You should optimize the multiplication process to avoid performing operations on zero values.

Function Signature

def multiply(A: list[list[int]], B: list[list[int]]) -> list[list[int]]:

Requirements

  • Let A be an m x k matrix and B be a k x n matrix. The result C will be an m x n matrix.
  • Return a dense representation of the resulting matrix C.
  • Your solution should only iterate over non-zero elements to perform multiplications.

Examples

Example 1: Input: A = [[1, 0, 0], [-1, 0, 3]] B = [[7, 0, 0], [0, 0, 0], [0, 0, 1]] Output: [[7, 0, 0], [-7, 0, 3]] Explanation: Only non-zero entries are multiplied. For instance, C[1][2] is computed as A[1][0]*B[0][2] + A[1][1]*B[1][2] + A[1][2]*B[2][2], which is (-1*0) + (0*0) + (3*1) = 3.

Constraints

  • m, n, k are between 1 and 500.
  • Input values are integers.
  • Optimize for time complexity by preprocessing non-zero indices.

Edge Cases

  • Matrices with all zeros.
  • Matrices where k is very large but non-zero elements are extremely rare.
  • Single element matrices.

Hints

Hint 1

Instead of a triple-nested loop, first identify the coordinates of all non-zero elements in matrix A.

Hint 2

Store the non-zero elements of matrix A in a dictionary or a list of tuples to quickly skip zeros.

Hint 3

Only iterate through the rows of A that contain non-zero elements and multiply them by the corresponding rows in B where the column index matches.

Follow-up questions

  • How would your approach change if both matrices were provided in Coordinate Format (COO) or Compressed Sparse Row (CSR) format?
  • What is the space complexity of your solution, and can you optimize it if the resulting matrix C is also expected to be sparse?
  • How would you parallelize this computation if the matrices were too large to fit into memory?