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.

Answer 1

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

Answer 2

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