如何用Numpy优化4×4096矩阵列间距离计算的O(N²)代码？

你的问题核心是Python循环处理O(N²)规模的计算效率太低，Numpy的向量化运算和矩阵操作可以直接利用底层优化的BLAS/LAPACK库，把计算速度提升几个数量级。

优化思路 #

欧几里得距离的平方对于复数向量$\boldsymbol{a}$和$\boldsymbol{b}$可以展开为：
$$|\boldsymbol{a} - \boldsymbol{b}|^2 = |\boldsymbol{a}|^2 + |\boldsymbol{b}|^2 - 2\operatorname{Re}(\boldsymbol{a} \cdot \overline{\boldsymbol{b}})$$
其中：

• $|\boldsymbol{a}|^2$是向量$\boldsymbol{a}$的L2范数平方（各元素模长的平方和）
• $\boldsymbol{a} \cdot \overline{\boldsymbol{b}}$是$\boldsymbol{a}$与$\boldsymbol{b}$共轭的内积，取实部后得到距离平方公式里的交叉项

利用这个公式，我们可以通过矩阵运算批量计算所有列对的距离平方，再推导得到距离，完全避免Python层面的循环。

优化代码 #

import numpy as np

# 生成复数测试数据（替换成你的实际矩阵）
data = np.random.rand(4, 4096) + 1j * np.random.rand(4, 4096)

# 计算每列的L2范数平方
col_sq_norms = np.sum(np.abs(data)**2, axis=0)

# 计算所有列对的内积实部（data.T是4096×4，data.conj()是4×4096，点乘后得到4096×4096的内积矩阵）
inner_prods = np.dot(data.T, data.conj()).real

# 批量计算所有列对的距离平方
sq_dist_matrix = col_sq_norms[:, np.newaxis] + col_sq_norms - 2 * inner_prods

# 提取上三角部分（对应原代码中i ≤ j的列对，包含对角线）
upper_tri_dist_sq = np.triu(sq_dist_matrix)

# 计算所有距离并找最小值（处理浮点误差导致的极小负数）
all_norms = np.sqrt(np.maximum(upper_tri_dist_sq, 0))
min_norm = all_norms.min()

print(min_norm)

为什么更快？ #

• 所有核心计算都是Numpy的底层C实现，比Python循环的效率高得多
• 矩阵乘法和广播操作充分利用了CPU的缓存和并行优化，尤其适合大规模数据
• 避免了原代码中重复的向量减法和范数计算，减少了冗余操作

如果不需要包含i=j的情况（即排除距离为0的对角线元素），可以修改最后两步：

# 提取上三角且排除对角线（i < j）
upper_tri_dist_sq = np.triu(sq_dist_matrix, k=1)
# 只取非零元素计算距离
all_norms = np.sqrt(upper_tri_dist_sq[upper_tri_dist_sq > 0])
min_norm = all_norms.min()

内容的提问来源于stack exchange，提问作者TolgaT