python - 有没有办法让这段代码更通用、更优雅？

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)

或者

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

None 在放置在下标中的位置创建一个大小为 1 的新轴。省略号 (...) 会获取剩余的轴，因此您不必显式编写 :, :。

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