numpy数组整数索引在多个维度

Question

我很确定我错过了整数索引的东西,可以使用一些帮助.假设我创建了一个2D数组:

>>> import numpy as np
>>> x=np.array(range(24)).reshape((4,6))
>>> x
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23]])

然后,我可以选择第1行和第2行:

>>> x[[1,2],:]
array([[ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17]])

或者第2行和第3行的第1列:

>>> x[[1,2],1]
array([ 7, 13])

因此,对我来说,我可以选择第1行和第2行的第3,4和5列:

>>> x[[1,2],[3,4,5]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: shape mismatch: objects cannot be broadcast to a single shape

相反,我需要分两步完成:

>>> a=x[[1,2],:]
>>> a
array([[ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17]])
>>> a[:,[3,4,5]]
array([[ 9, 10, 11],
       [15, 16, 17]])

来自R,我的期望似乎是错误的.您能否确认一步确实无法做到这一点,或建议更好的选择？谢谢!

编辑:请注意我在示例中选择的行和列恰好是连续的,但它们不一定是.换句话说,切片索引不适用于我的情况.

Answer 1

您还可以选择在索引数组中使用广播,这是我通常会做的,而不是索引两次,这会创建数据的中间副本:

>>> x[[[1], [2]],[[3, 4, 5]]]
array([[ 9, 10, 11],
       [15, 16, 17]])

为了更好地了解正在发生的事情以及如何处理更多的指数:

>>> row_idx = np.array([1, 2])
>>> col_idx = np.array([3, 4, 5])
>>> x[row_idx.reshape(-1, 1), col_idx]
array([[ 9, 10, 11],
       [15, 16, 17]])