Efficient Numpy slicing on 3d array

Question

I am trying to find the most efficient way to do slicing for a 3D numpy array. This is a subset of the data, just for test purposes :

in_arr =np.array([[[0,1,2,5],[2,3,2,6],[0,1,3,2]],[[1,2,3,4],[3,1,0,5],[2,4,0,1]]])
indx =[[3,1,2],[2,0,1]]

I need to get the value at the indx as stated. For example, indx[0][0] is 3, so I am looking for the 3rd elem of in_arr[0][0], in this case, 5.

I have the following code that will do what i need it to do, but the time complexeity is n^2, which I am not happy about.

list_in =[]
for x in range(len(indx)):
    arr2 = []
    for y in range(len(indx[x])):
        arr2.append(in_arr[x][y][indx[x][y]])
        #print in_arr[x][y][indx[x][y]]
    list_in.append(arr2)
print list_in

I am looking for a very fast and efficient way to do the same task for a large dataset.

Indexing docs For a real answer, improve your example a bit. E.g. by adding a complete expected output for your task. — sascha
– sascha, Commented Dec 8, 2017 at 16:19

jakevdp · Accepted Answer · 2017-12-08 16:26:06Z

1

You can do this efficiently using broadcasted arrays of indices; for example:

i1 = np.arange(2)[:, np.newaxis]
i2 = np.arange(3)[np.newaxis, :]
i3 = np.array(indx)
in_arr[i1, i2, i3]
# array([[5, 3, 3],
#        [3, 3, 4]])

What numpy does here is to effectively match the entries of the three index arrays, and extract the associated entries from in_arr: the reason for the [:, np.newaxis] and [np.newaxis, :] terms is that it reshapes the three arrays to be compatible via numpy's broadcasting rules.

answered Dec 8, 2017 at 16:26

jakevdp

88.6k15 gold badges140 silver badges177 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Martin Over a year ago

yes, i was looking at this and I am now understanding it. Really appreciate your help.

Collectives™ on Stack Overflow

Efficient Numpy slicing on 3d array

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related