如何在Numpy中无需双层循环实现3D矩阵的逐数组比较

嘿，完全理解你不想用双层循环来做这个——毕竟NumPy的优势就是向量化操作，能帮我们摆脱低效的Python循环！

你要实现的是对三维数组中每个二维子数组两两比较是否完全相等，生成一个对称的布尔真值表。这里有两种高效的向量化方案：

方案1：利用广播与all()方法 #

这是最直观的向量化思路，借助NumPy的广播机制来实现两两比较：

import numpy as np

M = np.array([[[1,2,3], [1,3,4]], [[3,4,5], [1,2,3]], [[1,2,3], [1,3,4]]])

# 扩展维度实现广播，然后在子数组的所有维度上判断全相等
result = (M[:, None] == M).all(axis=(2, 3))

原理说明： #

• M[:, None]将原数组从形状(3,2,3)变为(3,1,2,3)，这样和原数组M广播后会得到一个(3,3,2,3)的四维数组，其中每个位置对应的是M[i]和M[j]逐元素比较的结果。
• .all(axis=(2,3))会在第2和第3维度（也就是每个子数组的内部维度）上判断所有元素是否都相等，最终压缩成(3,3)的布尔矩阵。

方案2：利用视图打包子数组（更高效） #

如果你的子数组尺寸很大，这个方法会更快——它把每个二维子数组打包成一个单一的“虚拟元素”，然后只需要比较这些虚拟元素即可：

import numpy as np

M = np.array([[[1,2,3], [1,3,4]], [[3,4,5], [1,2,3]], [[1,2,3], [1,3,4]]])

# 计算每个子数组对应的字节大小，生成void类型的视图
subarray_bytes = M.dtype.itemsize * M.shape[1] * M.shape[2]
M_views = M.view(np.dtype((np.void, subarray_bytes)))

# 广播比较视图，得到布尔矩阵
result = (M_views[:, None] == M_views).astype(bool)

原理说明： #

• view方法不会复制数据，只是改变了数组的解读方式：把每个(2,3)的子数组看成一个连续的字节块，用np.void类型表示，这样M_views就变成了形状(3,)的一维数组。
• 之后的广播比较就变成了一维数组的两两相等判断，效率比逐元素比较高很多。

验证结果 #

两种方案得到的结果都和你用双层循环生成的一致：

print(result)
# 输出：
# [[ True False  True]
#  [False  True False]
#  [ True False  True]]

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