如何返回NumPy数组中元素总和最高的(row, column)位置索引

你原思路的核心问题有两个：

• return语句缩进位置错误，放在了行循环的内部，第一轮行遍历结束就会直接返回，没有计算所有行列的总和
• 直接修改输入的原数组属于不良实践，容易导致其他依赖该数组的代码出现逻辑错误

最优实现（NumPy向量化写法，无需循环） #

直接利用NumPy的内置求和、索引接口实现，性能远高于手动循环：

import numpy as np

def find_max_sum_pos(array):
    # 对最后一维求和，得到每个(row, column)位置的总和数组，shape为(行数, 列数)
    sum_arr = np.sum(array, axis=-1)
    # 拉平索引转多维索引，得到最大值对应的行列坐标
    return np.unravel_index(np.argmax(sum_arr), sum_arr.shape)

测试效果：

array = np.array([[[7.1, 4.5, 2.1], [0.9, 0.6, 10.1]],
                  [[11.0, 5.4, 4.3], [6.7, 0.3, 8.2]]])
print(find_max_sum_pos(array))
# 输出：(1, 0)

原循环思路修正版 #

如果坚持用手动循环的思路，修正后的代码如下：

def function(array):
    max_sum = -float('inf')
    max_pos = (0, 0)
    for row in range(array.shape[0]):
        for col in range(array.shape[1]):
            current_sum = np.sum(array[row, col])
            if current_sum > max_sum:
                max_sum = current_sum
                max_pos = (row, col)
    return max_pos

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