Python NumPy多维数组按索引取值结果不符预期如何解决

问题原因 #

你当前的写法A[0,i]仅指定了数组A第0轴的索引，NumPy会将索引数组i的所有值作为第1轴的行索引，默认提取对应行的全部列元素，因此返回结果多了一整维的行内容，最终输出维度为(1,3,2,3)，和预期不符。
另外你对索引数组i的维度描述有误，实际定义的i维度为(1,3,2)，其中每一组长度为2的数组存储的是目标元素的行坐标、列坐标对，分别对应A[0]这个3×3矩阵的第1轴、第2轴位置。你需要提取的三个目标元素坐标分别是：

• (0,0) → 值1
• (1,2) → 值6
• (2,2) → 值9

修正代码 #

需要将索引数组拆分为独立的行索引、列索引，同时传入两个轴的索引做高级索引，最后补充长度为1的尾维度匹配预期输出格式：

import numpy as np
i = np.array([[[0,0],[1,2],[2,2]]])
A = np.array([[[1,2,3],[4,5,6],[7,8,9]]], dtype=float)

# 拆分坐标对为行索引、列索引
row_idx = i[..., 0]
col_idx = i[..., 1]
# 双轴高级索引取对应元素，新增尾维度匹配输出结构
I = A[0, row_idx, col_idx][..., np.newaxis]
print([I])

运行结果 #

执行代码后将输出你预期的结果：

[array([[[1.],
        [6.],
        [9.]]])]

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