python - 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])

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

现在，显然可以编写一个循环来从 foo 计算 bar，但这会很慢。还有比这更聪明的事吗？

更新 pd.Series 解决方案比纯 numpy 解决方案慢大约 7 倍。愚蠢的循环解决方案非常慢(不足为奇)，但是当使用 numba 编译的 jit 与 numpy 解决方案一样快。

Answer 1

你可以用pandas做cumcount

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)