Python numpy索引多维数组

Question

对于y的每一行，我想获取其索引在m中指定的元素。你知道吗

>>> y = np.arange(15).reshape(3,5)
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]]

>>> m = np.array([[0, 1], [1, 2], [2, 3]])

预期输出：

[[0, 1]
 [6, 7]
 [12, 13]]

具有for循环的溶液

>>> np.stack([y[i, cols] for i, cols in enumerate(m)])

有没有办法不用for循环就能做到这一点？你知道吗

Answer 1

有一行，虽然不比使用for循环的您自己的建议好多少：

y[..., m][np.identity(3, dtype=bool)]

虽然它会给你一些关于numpy索引的见解。你知道吗

Answer 2

使用一个数组中的值作为另一个数组的索引称为“花式索引”，但是索引操作将对所有行重复：

y = numpy.arange(15).reshape(3,5)
y[:, [0, 2, 3]]
# array([[ 0,  2,  3],
#        [ 5,  7,  8],
#        [10, 12, 13]])

如果要单独“每行使用一个索引值”，则需要将该行与索引的关系作为另一个索引：

y[[0, 1, 2], [0, 2, 3]]
# array([ 0,  7, 13])

由于索引数组m是2D，因此需要告诉NumPy m中这两个维度中的哪个对应于行索引。通过向升序索引中添加另一个空轴（关键字：broadcasting）可以实现这一点，然后

y = numpy.arange(15).reshape(3,5)
m = numpy.array([[0, 1], [1, 2], [2, 3]])

y[numpy.arange(len(m))[:, None], m]
# array([[ 0,  1],
#        [ 6,  7],
#        [12, 13]])