Numpy: объединение / расширение 2D-массивов на основе совпадений в столбце. Как устранить за l oop?

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.tile inside a для `` l oop. Однако это неэффективно:

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 l oop и быстрее работает с большими наборами данных?

Answer 1

вы можете попробовать использовать столбец 0 из a для индексации 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]]