从NumPy矩阵提取指定索引值时返回矩阵而非向量的问题求助

嘿，我懂你遇到的问题了——你想用np.take提取矩阵每行中对应索引数组arr的元素，结果得到的是个矩阵而非你想要的一维向量，这确实有点坑。咱们来拆解原因，再给出最直接的解决方案。

为什么np.take会返回矩阵？ #

当你使用np.take(Mat, arr, axis=1)时，numpy会把arr里的所有索引应用到矩阵的每一行上。比如你的示例中arr长度是6，它会给每一行都提取第0、0、3、0、0、3个元素，最终得到一个和原矩阵行数相同、列数等于arr长度的矩阵，这显然不是你要的“第i行取第arr[i]个元素”的效果。

两种可行的解决方案 #

方法1：使用高级索引（推荐，简洁高效） #

这是numpy中处理这类“每行取对应列”需求的标准做法，直接通过行索引和列索引的配对来提取元素：

import numpy as np

# 你的示例数据
Mat = np.array([[0., 4., 4., 4.], [2., 8., 8., 2.], [7., 7., 7., 7.], [5., 5., 5., 5.], [7., 7., 7., 7.], [2., 6., 6., 6.]])
arr = np.array([0, 0, 3, 0, 0, 3])

# 生成行索引数组（0到行数-1），和arr的列索引配对
row_indices = np.arange(Mat.shape[0])
result = Mat[row_indices, arr]

print(result)
# 输出：[0. 2. 7. 5. 7. 6.]

这里np.arange(Mat.shape[0])生成了和矩阵行数一致的行索引数组，和arr的列索引数组一一对应，numpy会自动提取每个(行索引, 列索引)位置的元素，直接返回一维向量。

方法2：调整np.take的用法 #

如果你坚持想用np.take，可以把arr转换成列向量，让numpy为每一行单独取对应位置的索引，最后再展平成一维向量：

# 将arr转为列向量，再用take提取
result = np.take(Mat, arr.reshape(-1, 1), axis=1).flatten()

# 或者用np.newaxis实现列向量转换
# result = np.take(Mat, arr[:, np.newaxis], axis=1).flatten()

print(result)
# 输出同样是：[0. 2. 7. 5. 7. 6.]

arr.reshape(-1,1)把一维的arr变成了(n,1)的列矩阵，此时np.take在axis=1时，会为第i行取第arr[i]个元素，得到一个(n,1)的二维矩阵，最后用flatten()把它转成一维向量。

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