如何从NumPy数组中按指定列索引获取特定格式的多列

你可以直接将多个列索引放入**列表（或NumPy数组）**中进行索引，就能得到保持二维结构的结果，无需使用切片语法。具体实现如下：

示例代码 #

1. 基础多列索引示例 #

import numpy as np

array_data = np.array([[1, 2, 3],
                       [4, 5, 6],
                       [7, 8, 9]])

# 要选取的多个列索引，比如第2列和第3列（索引1和2）
column_indices = [1, 2]
selected_columns = array_data[:, column_indices]
print(selected_columns)

输出结果：

[[2 3]
 [5 6]
 [8 9]]

2. 结合Pandas列名转索引的场景 #

如果你的列索引来自Pandas数据框的指定列名，可以先通过df.columns.get_indexer()获取对应索引，再传入NumPy数组进行索引：

import numpy as np
import pandas as pd

# 示例Pandas数据框
df = pd.DataFrame([[1,2,3], [4,5,6], [7,8,9]], columns=['col_a', 'col_b', 'col_c'])

# 指定要选取的列名
target_columns = ['col_b', 'col_c']
# 获取列名对应的索引
column_indices = df.columns.get_indexer(target_columns)

# 从NumPy数组中选取对应列
array_data = df.to_numpy()
selected_columns = array_data[:, column_indices]
print(selected_columns)

输出结果同样是二维结构：

[[2 3]
 [5 6]
 [8 9]]

原理说明 #

当你使用单个整数索引（如[:,1]）时，NumPy会返回降维的一维数组；但使用列表/数组形式的多个索引（如[:,[1,2]]）时，NumPy会保留原数组的二维结构，正好符合你需要的格式。

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