带有重置功能的 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

使用for循环。设置一个从 1 开始的变量，并在每次遇到不同的数字时重置它。例如：

counter = 1;
outputArray = [];
for number in npArray:
    if number == -1:
        outputArray.append(counter)
        counter += 1
    elif number == 1:
        outputArray.append(0)
    else:
        outputArray.append(0)
        counter = 1
print(outputArray)

Answer 2

这是对@skyrider 代码的修复

npArray = [1,1,-1,-1,0,-1,1]
counter = 0
outputArray = []
for number in npArray:
    if number == -1:
        counter += 1
        outputArray.append(counter)
    elif number == 0:
        outputArray.append(0)
        counter = 0
    else:
        outputArray.append(counter)
print(outputArray)

Answer 3

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

是，使用 Series.cumsum and 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])

循序渐进

创建等于 0 时重置的伪组：

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

然后groupby these pseudo-groups and 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)

Answer 4

让我们先将您的 numpy 数组保存在一个变量中：

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

我定义了一个变量，count 来保存你关心的值，并将它设置为零。然后我定义一个列表来保存新元素。我们称它为 l。然后我迭代 a 的元素，在每次迭代中我将元素命名为 i。在每次迭代中，我实现逻辑：

如果i为-1，则增加counter
否则，如果 i 为 0，则重置 counter
否则什么也不做最后，我将 counter 附加到 l。最后，将 l 转换为一个 numpy 数组，out.

l = []
count = 0
for i in a:
    if i == -1:
        count+=1
    elif i==0: 
        count = 0
    l.append(count)
out = np.array(l)
out

Answer 5

对于此基准测试，我似乎比使用 numba 的 Pandas 解决方案获得了 10 倍的加速：

from numba import jit

inp1 = np.array([1,1,-1,-1,0,-1,1], dtype=int)
inp2 = np.random.randint(-1, 10, size=10**6)

@jit
def with_numba(arr):
  val = 0
  put = np.zeros_like(arr)
  for i in range(arr.size):
    if arr[i] == -1:
      val += 1
    elif arr[i] == 0:
      val = 0
    put[i] = val

  return put

def with_pandas(inp):
  s = pd.Series(inp)
  return s.eq(-1).groupby(s.eq(0).cumsum()).cumsum().to_numpy()
  
assert (with_numba(inp1) == with_pandas(inp1)).all()
assert (with_numba(inp2) == with_pandas(inp2)).all()

%timeit with_numba(inp2)
# 100 loops, best of 5: 4.57 ms per loop
%timeit with_pandas(inp2)
# 10 loops, best of 5: 46.3 ms per loop

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

Numpy array counter with reset

python

numpy

pandas

cumsum

循序渐进

时间