使用3D Numpy数组进行矢量化操作

Question

我正在研究一种算法来匹配两种对象（让我们说balls和buckets ）。 每个对象都建模为4D numpy数组，每种对象都分组在另一个数组中。 我的方法基于计算每对(ball, bucket)之间所有可能的差异(ball, bucket)并对差异应用相似性函数。

我试图避免for循环，因为速度与我所做的事情确实相关，所以我通过重塑初始数组之一，广播numpy操作并创建3D numpy数组（ diff_map ）来创建这些差异。 我没有找到任何关于此的好的教程，因此我想知道是否有更“合适的方法”来做到这一点。 如果可能的话，我也希望看到有关这种操作（多维重塑和广播）的任何良好参考。

我的代码：

import numpy as np

balls = np.random.rand(3,4)
buckets = np.random.rand(6,4)
buckets = buckets.reshape(len(buckets), 1, 4)
buckets
array([[[ 0.38382622,  0.27114067,  0.63856317,  0.51360638]],

   [[ 0.08709269,  0.21659216,  0.31148519,  0.99143705]],

   [[ 0.03659845,  0.78305241,  0.87699971,  0.78447545]],

   [[ 0.11652137,  0.49490129,  0.76382286,  0.90313785]],

   [[ 0.62681395,  0.10125169,  0.61131263,  0.15643676]],

   [[ 0.97072113,  0.56535597,  0.39471204,  0.24798229]]])

diff_map = balls-buckets
diff_map.shape
(6, 3, 4)

对于循环

根据要求，这是我要避免的for循环：

diff_map_for = np.zeros((len(buckets), len(balls), 4))
for i in range(len(buckets)):
    for j in range(len(balls)):
        diff_map_for[i, j] = buckets[i]-balls[j]

`可以肯定的是，让我们比较两个结果：

np.all(diff_map == diff_map_for)
True

Answer 1

这对您有用吗？

import numpy as np

balls = np.random.rand(3,4)
buckets = np.random.rand(6,4)

diff_map = buckets[:, np.newaxis, :] - balls[np.newaxis, :, :]
print(diff_map.shape)
# output: (6, 3, 4)

# ... compared to for loop
diff_map_for = np.zeros((len(buckets), len(balls), 4))
for i in range(len(buckets)):
    for j in range(len(balls)):
        diff_map_for[i, j] = buckets[i] - balls[j]

print(np.sum(diff_map - diff_map_for))
# output: 0.0