python numpy `np.take` with 2 dimensional array

Question

I'm trying to take a list of elements from an 2D numpy array with given list of coordinates and I want to avoid using loop. I saw that np.take works with 1D array but I can't make it work with 2D arrays.

Example:

a = np.array([[1,2,3], [4,5,6]])
print(a)
# [[1 2 3]
#  [4 5 6]]

np.take(a, [[1,2]])
# gives [2, 3] but I want just [6]

I want to avoid loop because I think that will be slower (I need speed). But if you can persuade me that a loop is as fast as an existing numpy function solution, then I can go for it.

Answer 1

If I understand it correctly, you have a list of coordinates like this:

coords = [[y0, x0], [y1, x1], ...]

To get the values of array a at these coordinates you need:

a[[y0, y1, ...], [x0, x1, ...]]

So a[coords] will not work. One way to do it is:

Y = [c[0] for c in coords]
X = [c[1] for c in coords]

or

Y = np.transpose(coords)[0]
X = np.transpose(coords)[1]

Then

a[Y, X]