import numpy as np
x = np.array([[1,2 ,3], [9,8,7]])
y = np.array([[2,1 ,0], [1,0,2]])
x[y]
Expected output:
array([[3,2,1], [8,9,7]])
If x and y were 1D arrays, then x[y] would work. So what is the numpy way or most pythonic or efficient way of doing this for 2D arrays?
You need to define the corresponding row indices.
One way is:
>>> x[np.arange(x.shape[0])[..., None], y]
array([[3, 2, 1],
[8, 9, 7]])
You can calculate the linear indices from y and then use those to extract specific elements from x, like so -
# Linear indices from y, using x's shape
lin_idx = y + np.arange(y.shape[0])[:,None]*x.shape[1]
# Use np.take to extract those indexed elements from x
out = np.take(x,lin_idx)
Sample run -
In [47]: x
Out[47]:
array([[1, 2, 3],
[9, 8, 7]])
In [48]: y
Out[48]:
array([[2, 1, 0],
[1, 0, 2]])
In [49]: lin_idx = y + np.arange(y.shape[0])[:,None]*x.shape[1]
In [50]: lin_idx # Compare this with y
Out[50]:
array([[2, 1, 0],
[4, 3, 5]])
In [51]: np.take(x,lin_idx)
Out[51]:
array([[3, 2, 1],
[8, 9, 7]])
