Insert more than one value per row at index

user_13259978 0
Viewed 54

I am trying to vectorize the following operation:

  • Place a smaller array into a bigger array, whereby the index changes as a function of another array for each row.

Example data:

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

Desired output: array_combined = [[0,1,2,0,0],[3,4,0,0,0],[0,0,5,6,0]]

So far I have been getting it to work with apply_along_axis - but I am wondering if there is a more efficient way of solving the problem. I can't seem to wrap my head around the indexing necessary to solve the problem.

vectorization numpy python
1 Answers

We can use advanced indexing here:

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

Or better, generalizing for whatever shape array_small may have:

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]])
user_96986840
Related