在NumPy中索引行和列

Question

a = np.array(list(range(16).reshape((4,4))
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

假设我想要中间的方块。可能会这样做：

a[[1,2],[1,2]]

但是我得到了这个：

array([5, 10])

这个方法可以正常运行，但看起来不够优雅：

a[[1,2],:][:,[1,2]]
array([[5, 6],
       [9, 10]])

所以我的问题是：

1. 为什么是这样的？实现该方式所需的前提条件是什么？
2. 是否有一种规范的方法可以同时沿多个索引选择？

- generic_user

---

当使用列表（或实际上是数组）进行索引时，数组会互相“广播”。 x [[1,2]，[1,2]] 选择对角线，即点（1,1）和（2,2）。 x [[ [1]，[2]]，[1,2]] 选择块。在MATLAB中，索引块更简单，但“对角线”更棘手。 - hpaulj

Answer 1

你可以同时进行两个切片操作，而不是创建一个视图并再次进行索引：

import numpy as np

a = np.arange(16).reshape((4, 4))

# preferred if possible
print(a[1:3, 1:3])
# [[ 5  6]
#  [ 9 10]]

# otherwise add a second dimension to the first index to make it broadcastable
index1 = np.asarray([1, 2])
index2 = np.asarray([1, 2])
print(a[index1[:, None], index2])
# [[ 5  6]
#  [ 9 10]]

- Stefan B

---

对，我的真实用例涉及使用整数数组进行索引，就像问题中描述的那样。我已经编辑了我的问题以反映这一点。 - generic_user

抱歉，我误解了你的问题。我会更新答案。 - Stefan B

Answer 2

我觉得你可以在 高级索引 中读到更多细节。基本上，当您通过列表/数组切片数组时，这些数组将被广播并一起迭代。

在你的情况下，你可以这样做：

idx = np.array([1,3])
a[idx,idx[:,None]]

或者如上文档所述：

a[np.ix_(idx, idx)]

输出：

array([[ 5, 13],
       [ 7, 15]])

- Quang Hoang

Answer 3

你可以使用多个np.take从多个轴中选择索引

a = np.arange(16).reshape((4, 4))
idx = np.array([1,2])
np.take(np.take(a, idx, axis=1), idx, axis=0)

或者（稍微更易读）

a.take(idx, axis=1).take(idx, axis=0)

输出：

array([[ 5,  6],
       [ 9, 10]])

np.take 还允许您方便地处理越界索引等问题。

- Kevin