Katkat
Back

Implement Sparse Matrix Operations

MetaMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenCoding & Algorithms

Last updated

This question tests the ability to design an efficient data structure for sparse matrix storage and implement optimized multiplication between a sparse matrix and a dense vector.

Sparse Matrix-Vector Multiplication

In machine learning, we often encounter matrices where most elements are zero. Storing these in a standard 2D array is memory-inefficient. Implement a class SparseMatrix that stores a matrix using the Compressed Sparse Row (CSR) format and supports efficient vector multiplication.

Requirements

  1. __init__(self, matrix): Initialize the sparse structure from a standard 2D list of integers. You must convert it to three arrays: values (non-zero elements), col_indices (column index for each non-zero element), and row_ptr (the starting index in values for each row).
  2. multiply(self, vector): Given a dense 1D list vector of length equal to the matrix columns, return the resulting dense 1D list of the matrix-vector product.

Constraints

  • The matrix dimensions are up to 10,000 x 10,000.
  • The number of non-zero elements is at most 1,000,000.
  • The vector length matches the number of columns.

Example

# Matrix: [[1, 0, 2], [0, 0, 3]]
# Vector: [1, 2, 3]

# Sparse representation:
# values: [1, 2, 3]
# col_indices: [0, 2, 2]
# row_ptr: [0, 2, 3]

# Result: [1*1 + 2*3, 3*3] -> [7, 9]

Edge Cases

  • A matrix with all zeros.
  • A matrix with a single non-zero element.
  • An empty matrix (0x0).

Hints

Hint 1

Think of `row_ptr` as a mapping that tells you exactly which slice of the `values` array corresponds to a specific row.

Hint 2

For each row `i`, iterate from `row_ptr[i]` to `row_ptr[i+1]` to identify the non-zero values and their corresponding column indices.

Hint 3

The result index `i` is the dot product of the non-zero segments found in the sparse row and the corresponding values in the input vector.

Follow-up questions

  • How would you modify your implementation to support Sparse Matrix-Sparse Matrix multiplication?
  • What are the space complexity trade-offs between CSR and Coordinate List (COO) formats?
  • If the matrix is updated frequently (elements changing from zero to non-zero), is CSR still the optimal choice? Why or why not?