Reputation: 287
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.
Upvotes: 1
Views: 1748
Reputation: 42776
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])])
Upvotes: 0
Reputation: 221674
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]])
Upvotes: 3