Numpy：根据一列中的匹配来合并/扩展2D数组。如何消除for循环？

Question

我正在寻找一种更“ Numpy”的方式来执行一项操作，以重复的术语将数据从一个数组扩展和连接到另一个数组。

示例数据

我希望a和b最终看起来像c：

a = np.array(((0, 13), (0, 14), (1, 15), (1, 16), (2, 17)))
b = np.array(((0, 415, 666), (1, 286, 583), (2, 777, 32), (3, 663, 12)))

c = np.array(((0, 13, 415, 666), (0, 14, 415, 666), (1, 15, 286, 583), (1, 16, 286, 583), (2, 17, 777, 32)))

在我的示例数据中，请注意，第一列是数组a中的某些元素的具有多个重复的排序的“索引”。还请注意，b包含未使用的索引。

工作代码的MWE

我目前正在使用np.repeat和np.tileinside afor``循环的组合来执行此任务。但是，这效率低下：

c = np.empty((0, 4))
for i in np.unique(a[:, 0]):
    foo = a[np.where(a[:, 0] == i)]
    bar = b[np.where(b[:, 0] == i)]
    output = np.vstack((output, np.column_stack((np.repeat(foo, np.size(bar, 0), 0), np.tile(bar, (np.size(foo, 0), 1))))))

问题

是否有一种方法可以使用纯Numpy执行此任务，从而消除for循环并在大型数据集上更快地执行？

Answer 1

您可以尝试使用a的第0列为b值编制索引（对不起，此处不确定该术语），然后使用hstack，例如：

c = np.hstack([a, b[a[:, 0], 1:]])
print (c)
[[  0  13 415 666]
 [  0  14 415 666]
 [  1  15 286 583]
 [  1  16 286 583]
 [  2  17 777  32]]