Sparse matrix multiplication

Background

Many real matrices — graph adjacency matrices, one-hot encodings, term-document matrices — are mostly zeros. Sparse matrix multiplication computes the product while skipping the zero entries, so the cost scales with the number of non-zeros instead of the full $O(mkn)$ of dense multiplication. It is a core primitive in graph algorithms, NLP pipelines, and recommender systems.

Problem statement

Implement sparse_matmul(A, B) that returns the product $C = AB$, computed by iterating only over non-zero entries. For $A$ of shape $(m, k)$ and $B$ of shape $(k, n)$:

$$C_{ij} = \sum_{\ell=1}^{k} A_{i\ell}\, B_{\ell j}$$

Skip any term where $A_{i\ell} = 0$ (and, ideally, where $B_{\ell j} = 0$).

Input

• A — list[list[float]] of shape (m, k).
• B — list[list[float]] of shape (k, n). (A's column count equals B's row count.)

Output

Returns list[list[float]] of shape (m, n): the product $AB$.

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 the non-zero products contribute — $C_{00} = 1\cdot 7 = 7$, $C_{10} = -1\cdot 7 = -7$, and $C_{12} = 3\cdot 1 = 3$; every other entry stays 0.

Constraints

• Skip multiplications where the left factor $A_{i\ell}$ is zero (an inner skip on $B_{\ell j}$ is a further optimisation).
• The result has shape (m, n); initialise it with zeros.
• The numeric result must equal the dense product A @ B.

Notes

• The speedup comes from touching only non-zeros: for density $\rho$ the work is roughly $O(\rho\, mkn)$ instead of $O(mkn)$.
• Production libraries store sparse matrices in CSR/CSC formats, which make enumerating "only the non-zeros" an $O(1)$-per-element operation.