numpy get values in array of arrays of arrays for array of indices

Question

I have a np array of arrays of arrays:

arr1 = np.array([[1,2,3],[4,5,6],[7,8,9]])
arr2 = np.array([[10,20,30],[40,50,60],[70,80,90]])
arr3 = np.array([[15,25,35],[45,55,65],[75,85,95]])

list_arr = np.array([arr1,arr2,arr3])

and indices array:

indices_array = np.array([1,0,2])

I want to get the array at index 1 for the first (array of arrays), the array at index 0 for the second (array of arrays) and the array at index 2 for the third (array of arrays)

expected output:

#[[ 4  5  6]
#[10 20 30]
#[75 85 95]]

I am looking for a numpy way to do it. As I have large arrays, I prefer not to use comprehension lists.

Divakar · Accepted Answer · 2015-11-19 09:47:34Z

3

Basically, you are selecting the second axis elements with indices_array corresponding to each position along the first axis for all the elements along the third axis. As such, you can do -

list_arr[np.arange(list_arr.shape[0]),indices_array,:]

Sample run -

In [16]: list_arr
Out[16]: 
array([[[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9]],

       [[10, 20, 30],
        [40, 50, 60],
        [70, 80, 90]],

       [[15, 25, 35],
        [45, 55, 65],
        [75, 85, 95]]])

In [17]: indices_array
Out[17]: array([1, 0, 2])

In [18]: list_arr[np.arange(list_arr.shape[0]),indices_array,:]
Out[18]: 
array([[ 4,  5,  6],
       [10, 20, 30],
       [75, 85, 95]])

answered Nov 19, 2015 at 9:47

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

Comments

Netwave · Accepted Answer · 2015-11-19 09:47:43Z

0

Just acces by linking postions to desired indexes (0-1, 1-0, 2-2) as follows:

desired_array = np.array([list_arrr[x][y] for x,y in enumerate([1,0,2])])

answered Nov 19, 2015 at 9:47

Netwave

43.1k6 gold badges69 silver badges110 bronze badges

Collectives™ on Stack Overflow

numpy get values in array of arrays of arrays for array of indices

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related