python - 用于选择二维数组元素的矢量化代码

Question

下面的代码显然返回矩阵中 ind 位置的元素:

def select_coord(a,ind):
    return a[ind]

但是我不知道如何矢量化它。换句话说:

b=np.asarray([[2,3,4,5],[7,6,8,10]])
indices=np.asarray([2,3])
select_coord(b,indices)

应返回[4,10]。

可以用for循环编写:

def new_select_record(a,indices):
    ret=[]
    for i in range a.shape[0]:
        ret.append(a[indices[i]])
    return np.asarray(ret)

有没有办法以矢量化的方式编写它？

Answer 1

获取b[0, 2]、b[1, 3]:

>>> import numpy as np
>>> b = np.array([[2,3,4,5], [7,6,8,10]])
>>> indices = np.array([2, 3])
>>> b[np.arange(len(indices)), indices]
array([ 4, 10])