Pinterest 稀疏矩阵：COO 存储、加法与乘法

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练版矩阵 shape 为 (rows,columns)，输入为 COO 三元组 (row,column,value)，同坐标可重复并求和；全用整数计算以便精确验证。
• 坐标越界报错；加法 shape 必须相同，乘法 A 的列数等于 B 的行数；结果按行列排序且不保存零值。

解题路线

1. 先校验 shape 和坐标，将每个矩阵的重复 COO 坐标合并，合并为零的项删除。
2. 加法对 A 的坐标字典复制一份，逐个加入 B 的值，再删除抵消为零的坐标。
3. 乘法把 B 的非零项按行 k 建索引；遍历每个 A(i,k)，只与 B(k,j) 的非零项相乘并累加到结果 (i,j)。
4. 对稠密小矩阵构造独立三重循环基线，随机覆盖重复坐标、抵消、空维度和形状不匹配。

手动走一遍

A 为 2×2，非零 (0,1,2)、(1,0,3)；B 为 2×2，非零 (1,0,4)、(0,1,5)。A+B 有 (0,1,7) 和 (1,0,7)；A×B 有 (0,0)=8、(1,1)=15。若 A 又有 (0,1,-2)，该坐标先抵消为零。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

"""COO training interface for sparse matrix storage, addition and multiplication."""

from collections import defaultdict


def _normalize(shape: tuple[int, int], entries):
    rows, columns = shape
    if rows < 0 or columns < 0:
        raise ValueError("negative matrix dimension")
    result = defaultdict(int)
    for row, column, value in entries:
        if not (0 <= row < rows and 0 <= column < columns):
            raise ValueError("COO coordinate out of bounds")
        result[(row, column)] += value
    return {position: value for position, value in result.items() if value != 0}


def add_sparse(shape_a, entries_a, shape_b, entries_b):
    if shape_a != shape_b:
        raise ValueError("addition requires equal shapes")
    result = defaultdict(int, _normalize(shape_a, entries_a))
    for position, value in _normalize(shape_b, entries_b).items():
        result[position] += value
    return sorted((row, column, value) for (row, column), value in result.items() if value != 0)


def multiply_sparse(shape_a, entries_a, shape_b, entries_b):
    if shape_a[1] != shape_b[0]:
        raise ValueError("inner dimensions must match")
    left = _normalize(shape_a, entries_a)
    right_rows = defaultdict(list)
    for (row, column), value in _normalize(shape_b, entries_b).items():
        right_rows[row].append((column, value))
    result = defaultdict(int)
    for (row, shared), left_value in left.items():
        for column, right_value in right_rows[shared]:
            result[(row, column)] += left_value * right_value
    return sorted((row, column, value) for (row, column), value in result.items() if value != 0)

复杂度

设 A/B 非零项数 a、b，B 第 k 行非零数 b_k，乘法候选对 P=Σ_(i,k)∈A b_k；预处理 O(a+b)，乘法累加 O(P)，输出排序 O(z log z)，z 为非零结果数；额外空间 O(a+b+z)。加法合并 O(a+b) 加结果排序。

易错点与追问

容易失分的边界

• 对重复 COO 坐标覆盖而非求和。
• 乘法把 B 的列号错当共享维 k。
• 结果相加抵消为零后仍声称是非零项。
• 把 COO 训练格式写成题目指定格式；题目未明确。

继续深挖

• CSR/CSC 对矩阵乘法的访问模式有什么好处？
• 两矩阵稀疏度很不均匀时，如何选择索引方向？
• 如果数值是浮点数，怎样定义零阈值和可复现累加？