python - 带有重置功能的 Numpy 数组计数器

Question

我有一个只有 -1、1 和 0 的 numpy 数组，如下所示:

np.array([1,1,-1,-1,0,-1,1])

我想要一个新数组来计算遇到的 -1。计数器必须在出现 0 时重置，出现 1 时保持不变:

期望的输出:

np.array([0,0,1,2,0,1,1])

当与更大的数组(最多 100 000)一起使用时，该解决方案必须非常少耗时

---

编辑:感谢您的贡献，我现在有一个可行的解决方案。

我仍在寻找一种非迭代的方法来解决它(没有 for 循环)。也许使用 Pandas 系列和 cumsum() 方法？

Answer 1

Maybe with a pandas Series and the cumsum() method?

是的，使用 Series.cumsum和 Series.groupby :

s = pd.Series([1, 1, -1, -1, 0, -1, 1])

s.eq(-1).groupby(s.eq(0).cumsum()).cumsum().to_numpy()
# array([0, 0, 1, 2, 0, 1, 1])

---

循序渐进

1. 创建等于 0 时重置的伪组:

   groups = s.eq(0).cumsum()
   # array([0, 0, 0, 0, 1, 1, 1])
   

2. 然后 groupby这些伪组和 cumsum当等于-1时:

   s.eq(-1).groupby(groups).cumsum().to_numpy()
   # array([0, 0, 1, 2, 0, 1, 1])
   

---

时间

not time consuming when used with larger array (up to 100,000)

groupby + cumsum比循环快 8 倍，给定 np.random.choice([-1, 0, 1], size=100_000):

%timeit series_cumsum(a)
# 3.29 ms ± 721 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit miki_loop(a)
# 26.5 ms ± 925 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit skyrider_loop(a)
# 26.8 ms ± 1.36 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)