如何从一维数组最小值索引提取二维数组对应列？

想要找到一维数组最小值的索引，再提取二维数组的对应列。当前获取索引的代码如下：

find_index = np.where(min(function()))

其中function()会生成含8个值的一维数组，这一步看似正常，但提取8×8数组varr对应列时遇到问题，尝试过两种写法都失败：

find_column = varr[:,find_index]

column_needed = [:,find_index]
find_column = np.take(varr, column_needed)

简化示例代码：

varr = np.array([1],[2],[3]) 
varr2 = np.array([1,2,3], [4,5,6], [7,8,9]) 
find_index = np.where(min(varr)) 
find_column = varr2[:,find_index]

先修正几个关键错误，再给出可行实现：

• 数组创建错误：原示例里的数组创建语法不符合NumPy要求，一维数组应写为np.array([1,2,3])，二维数组需要外层嵌套方括号：np.array([[1,2,3], [4,5,6], [7,8,9]])。

• 获取最小值索引的错误：min(varr)返回的是单个标量值，np.where()需要传入布尔条件数组而非标量。正确写法分两种场景：

  • 单个最小值索引（最快方式）：用np.argmin()直接获取最小值的第一个索引：

    find_index = np.argmin(varr)
    

  • 多个最小值索引（当存在多个相同最小值时）：用np.where()匹配最小值并取出索引数组：

    min_val = np.min(varr)
    find_index = np.where(varr == min_val)[0]  # [0]用于取出索引数组
    

• 提取对应列：

  • 单个索引场景（np.argmin()返回标量）：直接索引即可：

    varr = np.array([1,2,3])
    varr2 = np.array([[1,2,3], [4,5,6], [7,8,9]])
    find_index = np.argmin(varr)  # 得到索引0
    find_column = varr2[:, find_index]
    # 结果：array([1, 4, 7])
    

  • 多个索引场景（np.where()返回数组）：索引后可通过squeeze()去掉多余维度：

    varr = np.array([1,1,3])
    varr2 = np.array([[1,2,3], [4,5,6], [7,8,9]])
    min_val = np.min(varr)
    find_index = np.where(varr == min_val)[0]  # 得到array([0, 1])
    find_column = varr2[:, find_index].squeeze()
    # 结果：array([[1, 2], [4, 5], [7, 8]])
    

另外，你尝试的np.take()用法有误，正确提取列的写法是np.take(varr2, find_index, axis=1)，效果和varr2[:, find_index]一致。

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