如何将最后一个索引显式指定给np.ufunc.reduceat还原

data = np.arange(6)

我想用 np.add.reduceat ¹

np.add.reduceat(data, [0, 6, 3])[::2]

我立刻得到一个错误

IndexError: index 6 out-of-bounds in add.reduceat [0, 6)

如果我这样做

np.add.reduceat(data, [0, 5, 3])[::2]

我得到了错误的答案(10应该是15):

array([10, 12])

我能想到的唯一解决办法是屏蔽需要最后一个索引的位置,从中减去1,然后再把最后一个元素加回去:

index = np.array([0, 6, 3])
mask = (index == data.size)
index[mask] -= 1
result = np.add.reduceat(data, index)
# Mask is shifted back by one because it's the previous element that needs to be updated
result[:-1][mask[1:]] += data[-1]

result[::2] 给出想要的答案。这看起来像是一个巨大的乱七八糟的东西,我希望是一个优雅的一行(比这更快)。

¹ 我是 意识到有更好的方法来做到这一点。为了说明的目的,这只是一个人为的例子。这个问题的真正问题源于试图解决这个问题 numpy: fast regularly-spaced average for large numbers of line segments / points