如何使用索引变量i访问维度未知的n维numpy数组A的指定元素

直接将索引列表i转换为元组后传入数组即可完成取值，写法如下：

result = A[tuple(i)]

原理说明 #

• A[x,y,z]这类多维度索引写法的本质是隐式把逗号分隔的索引值封装为元组传入，等价于A[(x,y,z)]
• 若直接传入列表i，numpy会将其判定为对数组第一个维度的花式索引，无法匹配多维度取值的需求
• 转成元组后可以自动适配任意维度的数组，只要i的长度和数组A的维度一致，就能正确取出对应位置的元素

示例验证 #

import numpy as np

# 2维数组测试
a2 = np.array([[1,2,3], [4,5,6], [7,8,9]])
i2 = [1, 2]
print(a2[tuple(i2)])
# 输出：6

# 4维数组测试
a4 = np.random.randn(2,3,4,5)
i4 = [0,2,1,3]
print(a4[tuple(i4)])
# 输出对应位置的随机数值，无报错

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