我有一个平面数组b
:
a = numpy.array([0, 1, 1, 2, 3, 1, 2])
以及一个c
索引数组,标记每个"块"的开头:
b = numpy.array([0, 4])
我知道我可以使用缩减在每个"块"中找到最大值:
m = numpy.maximum.reduceat(a,b)
>>> array([2, 3], dtype=int32)
但。。。有没有办法通过矢量化操作(无列表、循环)找到块</edit>
(如 numpy.argmax
)内最大<edit>
的索引?
借用this post
的想法。
涉及的步骤:
-
按限制偏移偏移组中的所有元素。对它们进行全局排序,从而限制每个组停留在其位置,但对每个组中的元素进行排序。
-
在排序数组中,我们将查找最后一个元素,即组 max。他们的索引将是抵消组长度后的 argmax。
因此,矢量化实现将是 -
def numpy_argmax_reduceat(a, b):
n = a.max()+1 # limit-offset
grp_count = np.append(b[1:] - b[:-1], a.size - b[-1])
shift = n*np.repeat(np.arange(grp_count.size), grp_count)
sortidx = (a+shift).argsort()
grp_shifted_argmax = np.append(b[1:],a.size)-1
return sortidx[grp_shifted_argmax] - b
作为一个小的调整,可能更快的调整,我们也可以使用 cumsum
创建shift
,从而获得早期方法的变体,如下所示 -
def numpy_argmax_reduceat_v2(a, b):
n = a.max()+1 # limit-offset
id_arr = np.zeros(a.size,dtype=int)
id_arr[b[1:]] = 1
shift = n*id_arr.cumsum()
sortidx = (a+shift).argsort()
grp_shifted_argmax = np.append(b[1:],a.size)-1
return sortidx[grp_shifted_argmax] - b