时间戳

首页 > 解决方案 > Python numpy索引多维数组

python - Python numpy索引多维数组

问题描述

对于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循环的情况下做到这一点？

标签： pythonnumpy

解决方案

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

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是二维的，您需要告诉 NumPy 这两个维度中的哪一个m对应于您的行索引。你可以通过在升序索引中添加另一个空轴来做到这一点，（关键字：广播），你得到

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

推荐阅读