从 O(n³) 到按需计算：Swift 玩转稀疏矩阵乘法

摘要

在大多数算法题里，矩阵乘法都不算太陌生了。但一旦题目提示“稀疏矩阵”——也就是大部分值都是 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 循环都值得执行。