Select elements in each row with column indices from another array

user_10166571 0
Viewed 494

As the title say, given a 2d numpy array, for instance

a = np.array([[2,3,4,5],[12,4,5,7],[14,2,5,6],[12,3,4,6]])
idx = np.array([[0,2],[2,3],[1,3],[1,3]])

I want to select the first and third elements, 2, 4 from the first row, etc. So the final answer should be

ans = [[2,4],
       [5,7],
       [2,6],
       [3,6]]

I have tried both np.choose and np.take but I believe np.take flattens out the array, and np.choose doesn't look like what I expected.

Any idea would be appreciated! Many thanks!

numpy-ndarray numpy python
3 Answers

using List comprehension, you can loop over the 2 numpy array and create a new list that have the desired output. Later, you take each 2 elements and add them to a numpy array, this can be achieved using the reshape function:

aa = [a[index][x] for index, value in enumerate(idx) for x in value]
# aa = [2, 4, 5, 7, 2, 6, 3, 6]
aa = np.reshape(aa, (-1, 2))
print(aa)
# output 
[[2 4]
 [5 7]
 [2 6]
 [3 6]]
user_98778270

You could define some complex function that turns your idx into slices, but I think it would be much easier to access the array via list comprehension and enumerate, and then turn the list into an array. Like this:

a = np.array([[2,3,4,5],[12,4,5,7],[14,2,5,6],[12,3,4,6]])
idx = np.array([[0,2],[2,3],[1,3],[1,3]])

np.array([a[row,elems] for row,elems in enumerate(idx)])

result:

array([[2, 4],
       [5, 7],
       [2, 6],
       [3, 6]])
user_134122140

You can use advance indexing in numpy: 

a[np.indices(idx.shape)[0], idx]

np.indices(idx.shape)[0] creates the corresponding row indices for your column indices in idx and together, they form advance indexing.

output: 

[[2 4]
 [5 7]
 [2 6]
 [3 6]]
user_49759810
Related