CODEWITHSUNDEEP
pythonnumpyvectorization
irozada
irozada

Reputation: 11

Vectorizing a numpy array call of varying indices

I have a 2D numpy array and a list of lists of indices for which I wish to compute the sum of the corresponding 1D vectors from the numpy array. This can be easily done through a for loop or via list comprehension, but I wonder if it's possible to vectorize it. With similar code I gain about 40x speedups from the vectorization.

Here's sample code:

import numpy as np
indices = [[1,2],[1,3],[2,0,3],[1]]
array_2d = np.array([[0.5, 1.5],[1.5,2.5],[2.5,3.5],[3.5,4.5]])
soln = [np.sum(array_2d[x], axis=-1) for x in indices]

(edit): Note that the indices are not (x,y) coordinates for array_2d, instead indices[0] = [1,2] represents the first and second vectors (rows) in array_2d. The number of elements of each list in indices can be variable.

This is what I would hope to be able to do:

vectorized_soln = np.sum(array_2d[indices[:]], axis=-1)

Does anybody know if there are any ways of achieving this?

Upvotes: 1

Views: 234

Answers (1)

Pablo
Pablo

Reputation: 2481

First to all, I think you have a typo in the third element of indices...

The easy way to do that is building a sub_array with two arrays of indices:

i = np.array([1,1,2])
j = np.array([2,3,?])
sub_arr2d = array_2d[i,j]

and finally, you can take the sum of sub_arr2d...

Upvotes: 1

Related Questions