使用Numpy获取数组分块最小值对应索引时结果偏移如何修复

核心错误原因 #

原自定义函数的索引计算逻辑存在多余的偏移操作：在得到分组最小值的原始索引后，额外执行了减分块起始索引b + 全局index的操作，而调用时传入的b是index+1，最终导致所有返回索引整体偏移了-1，和预期结果不匹配。

修复代码 #

仅需要修改numpy_argmin_reduceat函数的最后返回逻辑即可，修改后的完整代码如下：

import numpy as np
# indexes          0, 1, 2,3,4, 5, 6,7, 8, 9,10, 11, 12
L_list = np.array([10,20,30,0,0,-5,11,2,33, 4, 5, 68, 7])
index =  np.array([3,5,7,11])

# 最小值计算逻辑保持不变
Drawdown = np.minimum.reduceat(L_list,index+1)

# 修复后的argmin reduceat函数
def numpy_argmin_reduceat(a, b):
    n = a.max() + 1  # limit-offset
    id_arr = np.zeros(a.size,dtype=int)
    id_arr[b] = 1
    shift = n*id_arr.cumsum()
    sortidx = (a+shift).argsort()
    grp_shifted_argmin = b
    # 直接返回排序后得到的原始索引，移除多余的偏移计算
    return sortidx[grp_shifted_argmin]


min_idx =numpy_argmin_reduceat(L_list,index+1)
#打印函数
DR_val_index = np.array([np.around(Drawdown,1), min_idx])
DR_result = np.apply_along_axis(lambda x: print(f'Min Values: {x[0]} at index: {x[1]}'), 0, DR_val_index)

输出结果 #

运行后得到的输出和预期完全一致：

Min Values: -5.0 at index: 5
Min Values: 2.0 at index: 7
Min Values: 4.0 at index: 9
Min Values: 7.0 at index: 12

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