Numpy沿指定维度拆分多维数组 获取各子数组最大值的坐标

问题根源 #

你之前的实现出错的核心原因是reshape前没有调整目标维度的位置：Numpy默认按C顺序（行优先）存储数组，直接将原shape为(a, x, c...)的数组reshape为(x, -1)时，每行存储的并不是对应[:, i, :...]切片的所有元素，因此argmax得到的索引无法对应原数组坐标。

通用实现代码 #

以下是支持任意维度数组的向量化实现，无循环、符合Pythonic风格：

import numpy as np

def get_max_coords_along_axis(arr, target_axis):
    # 将目标维度移动到数组第0位
    arr_shifted = np.moveaxis(arr, target_axis, 0)
    # 提取除目标维度外的其余维度形状
    rest_dims_shape = list(arr.shape[:target_axis]) + list(arr.shape[target_axis+1:])
    # 展平为二维数组：每行对应目标维度一个切片的所有元素
    arr_flatten = arr_shifted.reshape(arr_shifted.shape[0], -1)
    # 求每个切片最大值的扁平化索引
    max_flat_indices = arr_flatten.argmax(axis=1)
    # 反解为其余维度的坐标
    return np.unravel_index(max_flat_indices, rest_dims_shape)

测试用例验证 #

代入你提供的测试数组：

testarray = np.array([[[0.17028444, 0.38504759, 0.64852725],
        [0.8344524 , 0.54964746, 0.86628204]],
       [[0.77089997, 0.25876277, 0.45092835],
        [0.6119848 , 0.10096425, 0.627054  ]],
       [[0.8466859 , 0.82011746, 0.51123959],
        [0.26681694, 0.12952723, 0.94956865]],
       [[0.28123628, 0.30465068, 0.29498136],
        [0.6624998 , 0.42748154, 0.83362323]]])

a_max, c_max = get_max_coords_along_axis(testarray, target_axis=1)
print(a_max) # 输出 [2 2]
print(c_max) # 输出 [0 2]

输出结果和预期完全一致：

• testarray[a_max[0], 0, c_max[0]] = testarray[2,0,0] = 0.8466859，为testarray[:,0,:]的最大值
• testarray[a_max[1], 1, c_max[1]] = testarray[2,1,2] = 0.94956865，为testarray[:,1,:]的最大值

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