Question

这是问题所在。假设我有一个矩阵A =

array([[ 1.,  0.,  2.],
       [ 0.,  0.,  2.],
       [ 0., -1.,  3.]])

和索引的矢量p = array([0, 2, 1])。我想将3x3矩阵A转换为长度为3的数组（称之为v），其中v [j] = A [j，p [j]]，j = 0,1,2。我可以通过以下方式实现：

v = map(lambda (row, idx): row[idx], zip(A, p))

因此对于上面的矩阵A和索引p的向量，我希望得到array([1, 2, -1])（即第0行的第0个元素，第1行的第2个元素，第2行的第1个元素）。

但是我可以通过使用本机numpy（即没有明确压缩然后映射）来实现相同的结果吗？感谢。

Answer 1

我不认为存在这样的功能。为了实现你想要的，我可以想到两个简单的方法。你可以这样做：

np.diag(A[:, p])

此处，数组p将作为每个行的列索引应用，以便在对角线上您将拥有您要查找的元素。

作为替代方案，您可以使用以下方法避免产生大量不必要的条目：

A[np.arange(A.shape[0]), p]