如何在Python中利用索引数组直接访问多维数组的指定元素？

嘿，这个问题其实是Numpy索引机制里的常见小陷阱，我来给你掰扯清楚～

首先咱们先复现你遇到的现象：

import numpy as np
A = np.arange(4).reshape(2,2)
a = [1,1]

# 你想要的是A[1,1]对应的3，但直接传a得到了重复的行
print(A[a])
# 输出：array([[2, 3], [2, 3]])

为什么会出现这个结果？ #

这是Numpy的索引规则导致的：当你传入一维列表/数组作为索引时，Numpy默认会把它当作第一个维度（行）的索引。这里a=[1,1]就表示“选取第1行，再选一次第1行”，所以得到了两行一模一样的[2,3]，形状变成了(2,2)。

而你真正想要的是“行索引1，列索引1”的多维标量索引，也就是需要用逗号分隔两个维度的索引值（A[1,1]）。

最简洁的解决方法 #

把你的索引列表a转换成元组就行：

print(A[tuple(a)])
# 输出：3

原理是什么？

在Python里，A[x,y]其实等价于A[(x,y)]——Python会自动把逗号分隔的参数打包成元组。当你传入元组(1,1)时，Numpy会把元组里的每个元素对应到不同的维度：第一个元素是行索引，第二个是列索引，刚好精准定位到你要的A[1,1]。

其他可选方式 #

如果你习惯用Numpy数组操作，也可以把a转成Numpy数组后拆分索引：

idx = np.array(a)
print(A[idx[0], idx[1]])
# 输出：3

不过这种方式和你原来写的A[a[0],a[1]]本质一样，只是把列表转成了Numpy数组，不如转元组来得简洁。

拓展：批量定位多个元素 #

要是你有一组索引，比如indices = [[0,0], [1,1], [0,1]]，想一次性获取这些位置的元素，可以把索引拆分成两个一维数组分别作为行和列的索引：

indices = np.array([[0,0], [1,1], [0,1]])
print(A[indices[:,0], indices[:,1]])
# 输出：array([0, 3, 1])

内容的提问来源于stack exchange，提问作者Marco