如何从给定的n维索引中获取值

Question

I have some index like this:

ind = [(1,5),(2,2),(3,1)]

and I have an array:

arr = np.arange(36).reshape(6,6)

and I would like to obtain the result:

arr[1,5], arr[2,2], arr[3,1]

What's the cleanest method?

I know I can do something like this:

c=np.empty(len(ind))
for i in len(ind):
    a,b = ind[i]
    c[i] = arr[a,b]

But I would like a more matrix-like method to deal with this problem, but not element-like method. any suggestion?

Answer 1

Why not just use a list comprehension:

ind = [(1,5),(2,2),(3,1)]
arr = np.arange(36).reshape(6,6)

result = [arr[i] for i in ind]

If you need to turn it into an numpy array you can just pass the resulting list to it:

result = np.array([arr[i] for i in ind])

Answer 2

you can write directly

c=[]
for i in range(len(ind)):
    c.append( arr[ ind[i][0], ind[i][1] ] )