Как я могу эффективно выполнять эту массивную манипуляцию?

Question

Мне дают массив, такой как:

a = numpy.array([1, 2, 3, 4],[1, 2, 3, 4],[1, 2, 3, 4])

Мне также предоставляется другой массив со всеми записями, большими или равными нулю и меньшими, чем размер a. Пример:

array_index = numpy.array([[0, 0, 1, 2], [0, 1, 2, 2], [0, 1, 1, 3]])

array_index в основном говорит, куда мне следует складывать элементы a. Например, первая запись array_index [0, 0, 1, 2] говорит о том, что элемент индекса 0 первой записи a должен оставаться прежним, второй индекс 0 говорит о том, что элемент индекса 1 a должен быть сложен в индекс 0, третий индекс 1 говорит, что элемент индекса 2 a должен быть наложен на индекс 1 и т. д. Теперь мне нужно создать другой массив такой же формы, как у my_array, чтобы записи соответствовали значению индекса a в array_index. В этом случае я бы сначала

array_desired = numpy.zeros(array_index.shape)

И я бы заполнил следующие значения:

array_desired = [[a[0]+a[1], a[2], a[3], 0], [a[0], a[1], a[2]+a[3], 0],
                 [a[0], a[1]+a[2], 0, a[3]]]

Для этого в качестве первого подхода я построил следующее для 0-го элемента:

stack_index = numpy.where(array_index == 0)
array_desired[stack_index] += a[stack_index]
print(array_desired)
>>> [[1 2 0 0] [1 0 0 0] [1 0 0 0]]

что-то, но не совсем то, что я хотел:

>>>[[3 0 0 0][1 0 0 0][1 0 0 0]] 

Есть идеи, как добиться правильной укладки по индексу, как я описал выше?

Обновление:

У меня есть метод, который работает для первого индекса:

temp_array = numpy.zeros(array_index.shape)
stack_index = numpy.where(array_index == 0)
temp_array[stack_index] = a[stack_index]
n = np.sum(temp_array, axis=(1))
array_desired[:,0]= n

Но этот метод все равно потребовал бы, чтобы я перебрал каждый индекс. Я хотел бы иметь что-то более эффективное.

Answer 1

Я думаю, это делает то, что вы хотите.

Вот ваши образцы данных:

In [98]: a
Out[98]: 
array([[1, 2, 3, 4],
       [1, 2, 3, 4],
       [1, 2, 3, 4]])

In [99]: array_index
Out[99]: 
array([[0, 0, 1, 2],
       [0, 1, 2, 2],
       [0, 1, 1, 3]])

Создайте b, чтобы сохранить результат, и сделайте сумму, используя np.add.at. row_index это просто массив 3х1 [[0], [1], [2]].

In [100]: b = np.zeros_like(a)

In [101]: row_index = np.arange(array_index.shape[0]).reshape(-1, 1)

In [102]: np.add.at(b, (row_index, array_index), a)

In [103]: b
Out[103]: 
array([[3, 3, 4, 0],
       [1, 2, 7, 0],
       [1, 5, 0, 4]])