python中的部分累积和

Question

假设我有一个numpy数组（或者pandas Series如果它使它更容易的话），它看起来像这样：

foo = np.array([1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0])

我想转成数组

bar = np.array([0, 1, 2, 3, 4,0, 1, 2, 0, 1, 2, 3])

其中条目是您需要向左走多少步才能找到1in foo。

现在，显然可以编写一个循环来计算bar，foo但这会非常慢。还有什么更聪明的事情可以做吗？

更新该pd.Series解决方案比纯解决方案慢大约 7 倍numpy。愚蠢的循环解决方案非常慢（毫不奇怪），但是当使用 jit 编译时，numba它与解决方案一样快numpy。

Answer 1

你可以cumcount这样做pandas

s = pd.Series(foo)
bar = s.groupby(s.cumsum()).cumcount().to_numpy()
Out[13]: array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3], dtype=int64)