用于选择2D阵列的元素的矢量化代码

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

怎么样：np.diag(b[:,[2,3]])？

Answer 2

获取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])