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]]