如何将最后一个索引明确指定给 np.ufunc.reduceat

Question

假设我有一个数组

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。

Answer 1

reduceat 我用的不多，但看起来你只能有一个开放式范围，一个add to the end。

一种解决方法是填充数组（是的，我通常反对使用np.append :)）：

In [165]: np.add.reduceat(np.append(x,0),[0,6,3])
Out[165]: array([15,  0, 12])

或具有完整的范围配对：

In [166]: np.add.reduceat(np.append(x,0),[0,6,3,6])
Out[166]: array([15,  0, 12,  0])

我省略了通常的 [::2] 以澄清发生了什么。