有没有办法更通用，更优雅地编写此代码？

Question

我是Python的新手，想问一下，是否有办法更优雅，更通用地编写此代码：

array = np.linalg.norm(np.array([
   X-mu[0],
   X-mu[1],
   X-mu[2],
   X-mu[3],
   X-mu[4],
]), axis=2)

X是RGB图像（NxD，300000x3阵列），mu是kxD（5x3）矩阵。最终，数组应为kxN或Nxk矩阵。

有没有一种方法可以遍历mu的索引而不是手动键入？

Answer 1

您正在尝试对(3k, 3)和(5, 3)数组执行减法运算。为了能够减少第二个轴以获得(3k, 5)或(5, 3k)的最终结果，您需要第一个尺寸一起广播。如果您引入新的坐标轴，这对于广播来说是微不足道的：

np.linalg.norm(X[:, None, :] - mu[None, ...], axis=-1)

OR

np.linalg.norm(X[None, ...] - mu[:, None, :], axis=-1)

None在下标放置的位置创建一个尺寸为1的新轴。省略号（...）会抓住其余的轴，因此您不必显式编写:, :。

以上表达式的差异分别广播到(3k, 5, 3)和(5, 3k, 3)。在这两种情况下，norm都会缩小最后一个轴，保留形状的前两个元素。根据您的内存布局，这两种方法可能会更快，有时会明显。如果这对您很重要，那么值得进行基准测试。