Вставить более одного значения в строку в индексе

Question

Я пытаюсь векторизовать следующую операцию:

• Поместить меньший массив в больший массив, в результате чего индекс изменяется как функция другого массива для каждой строки.

Пример данных:

array_large = [[0,0,0,0,0],[0,0,0,0,0],[0,0,0,0,0]] 
array_small = [[1,2],[3,4],[5,6]] 

array_index = [[1],[0],[2]] #*random index

Требуемый вывод: array_combined = [[0,1,2,0,0],[3,4,0,0,0],[0,0,5,6,0]]

До сих пор я получал его для работы с apply_along_axis - но мне интересно, если Существует более эффективный способ решения проблемы. Я не могу обернуться вокруг индексации, необходимой для решения проблемы.

Answer 1

Мы можем использовать расширенное индексирование здесь:

array_large[np.arange(array_large.shape[0])[:,None], array_index+[0,1]] = array_small

Или лучше, обобщение для любой формы array_small может иметь:

i = np.arange(array_large.shape[0])[:,None]
j = array_index+np.arange(array_small.shape[1])
array_large[i,j] = array_small

---

print(array_large)

array([[0, 1, 2, 0, 0],
       [3, 4, 0, 0, 0],
       [0, 0, 5, 6, 0]])