Filtering a multi dimensional array with indices using broadcasting

user_10639382 0
Viewed 175

I am wondering whether there is a way to filter a multidimensional array with an list of indexes.

For example, if there is a multidimensional array as so [[40,60,15],[25,30,45]], I want to obtain the resulting values when you filter with the following index [(0,2),(1,0)]. Note that the index is a list of tuples

Note : I am trying to find out the fastest method of doing this. I managed to do this using a for loop and map + lambda function, however was wondering if numpy broadcasting was a possibility?

 from numpy import np

 some_arr = np.array([40,60,15],[25,30,45])
 some_idx = [(0,2),(1,0)]
 
 #Method 1 : Using for loop
 new_list = []
 for idx in some_idx:
     new_list.append(some_arr[idx])

 # Method 2: Using lambda
 new_list = list(map(lambda idx : some_arr[idx], some_idx))

 # Method 3 : Using numpy ??? (i get an error ofcourse but this is what I tried to do)
 new_arr = some_arr[some_idx]

And I end up with an error for method 3 : index 2 is out of bounds for axis 0 with size 2 The output for this example should be [15,25]

numpy python
1 Answers

Use zip to instead have tuple of ((x_0, x_1...),(y_0,y_1...)) indices, which you can then use to perform integer array indexing along both axes:

some_arr = np.array([[40,60,15],[25,30,45]])
some_idx = [(0,2),(1,0)]

some_arr[tuple(zip(*some_idx))]
# array([15, 25])

Comparing with your first method:

new_list = []
for idx in some_idx:
    new_list.append(some_arr[idx])

np.array_equal(new_list, some_arr[tuple(zip(*some_idx))])
# True
user_96986840
Related