文章目录
摘要
在大多数算法题里，矩阵乘法都不算太陌生了。但一旦题目提示“稀疏矩阵”——也就是大部分值都是 0 的那种，这就提示我们：有优化空间。这篇文章就用 Swift 带大家一步步搞懂怎么写一个更高效的稀疏矩阵乘法逻辑，顺便聊聊背后的思路。
描述
我们手上有两个矩阵，A 和 B，想把它们乘起来。和普通乘法不同的是，这两个矩阵大多数位置上的值其实都是 0。这就像做一道题，本来有 100 步，结果有 80 步都是无用功——我们当然要跳过不做。
举个例子：
let A = [
[1, 0, 0],
[-1, 0, 3]
]
let B = [
[7, 0, 0],
[0, 0, 0],
[0, 0, 1]
]
如果直接用普通三重循环，确实能出结果，但性能很一般，特别是矩阵大了之后。我们的目标是让程序只做真正有意义的乘法。
解题思路
这里采用的思路叫“跳过零值优化”：
- 对于矩阵 A，我们逐行处理，只关注那些不为 0 的值。
- 每当发现 A[i][j] ≠ 0，就看看 B[j][k] 是否也不是 0；
- 如果两个都不是 0，我们才去执行乘法并更新结果矩阵。
这样一来，原来很多无意义的乘法（比如 0 * 99）就可以完全跳过了。
代码实现（Swift）
func multiply(_ A: [[Int]], _ B: [[Int]]) -> [[Int]] {
let m = A.count
let k = A[0].count
let n = B[0].count
var result = Array(repeating: Array(repeating: 0, count: n), count: m)
for i in 0..<m {
for j in 0..<k {
if A[i][j] != 0 {
for l in 0..<n {
if B[j][l] != 0 {
result[i][l] += A[i][j] * B[j][l]
}
}
}
}
}
return result
}
分析这个代码是怎么做的？
我们还是保留了三层循环结构，这是矩阵乘法本来的模样。只是我们在第二层和第三层分别加了判断：
- 只要 A[i][j] 是 0，我们就跳过；
- 就算 A[i][j] 不是 0，如果 B[j][l] 是 0，那也一样可以跳过。
这个判断看起来不起眼，但对稀疏矩阵非常关键，因为它会让我们真正只处理那些有用的格子，省下大量不必要的计算。
示例测试与输出结果
我们还是用上面的例子：
let A = [
[1, 0, 0],
[-1, 0, 3]
]
let B = [
[7, 0, 0],
[0, 0, 0],
[0, 0, 1]
]
let result = multiply(A, B)
print(result)
// 输出：[[7, 0, 0], [-7, 0, 3]]
解释一下：
- 第一行：A[0] 只有第一个元素是 1，对应 B 的第 0 行，只有 7 是非 0，所以最终结果是 [7, 0, 0]
- 第二行：A[1] 的第一个是 -1，第三个是 3，对应乘出来的是 [-7, 0, 3]
时间复杂度
如果什么优化都不做，三层循环就是 O(m × k × n)，也就是标准矩阵乘法复杂度。
但我们这里做了优化，只有当 A[i][j] 和 B[j][l] 都不为 0 时才会去执行一次乘法，因此在稀疏矩阵场景下，实际复杂度大大降低。
换句话说，最终时间消耗和矩阵中非 0 元素的个数直接相关。
空间复杂度
我们没有引入额外的大型结构，只是构建了一个和结果维度一样的矩阵用于输出。
所以空间复杂度就是 O(m × n)，跟输入的尺寸一样。
总结
这道题的精髓，其实不在“如何写矩阵乘法”，而是：
怎么让你只花力气在有价值的位置上。
在工程实践中，比如推荐系统、图神经网络、图像稀疏表示等，稀疏矩阵的优化处理也是基本功。用 Swift 实现这类高效处理思路，也能让我们在写代码时更接近底层逻辑。
下次再遇到这种结构，就要记得：不是所有的 for 循环都值得执行。