Python中如何按指定列表提取并打印NumPy数组的特定行？

问题说明 #

需要从Numpy数组A中提取J指定的行，当J=[[1,4]]时，目标是获取第1行和第4行组成的二维数组，但现有代码输出不符合预期。

错误原因 #

现有代码在循环中每次将A重新赋值为A[i]：

• 第一次循环后，A变为一维的第1行数组
• 第二次循环时，A[i]实际是取这个一维数组的第4个元素，导致输出错误

修正代码 #

直接利用Numpy的数组索引特性，一次性提取所有目标行：

import numpy as np

A = np.array([[1,2,3,4,5],
              [6,7,8,9,10],
              [11,12,13,14,15],
              [16,17,18,19,20],
              [21,22,23,24,25]])

J = [[1, 4]]

# 提取J[0]中的索引对应的行
result = A[J[0]]
print([result])

输出结果 #

运行修正后的代码，输出与期望一致：

[array([[ 6,  7,  8,  9, 10],
        [21, 22, 23, 24, 25]])]

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