Numpy sum运行非零值的累计长度

Question

寻找一个快速向量化的函数，返回连续非零值的滚动个数。每当遇到零时，计数应该从0开始重新计算。结果应与输入数组具有相同的形状。

给定这样的数组：

x = np.array([2.3, 1.2, 4.1 , 0.0, 0.0, 5.3, 0, 1.2, 3.1])

该函数应返回以下内容：

array([1, 2, 3, 0, 0, 1, 0, 1, 2])

- steve

Answer 1

本文列出了一种向量化方法，基本上包括两个步骤：

1. 初始化一个与输入向量x大小相同的零向量，并在对应于非零值的位置设置为1。

2. 接下来，在该向量中，我们需要在每个“岛屿”的结束/停止位置之后放置每个岛屿的负运行长度。这样做意图是稍后再次使用cumsum函数，这将导致“岛屿”中的连续数字和其他地方的0。

以下是实现代码 -

import numpy as np

#Append zeros at the start and end of input array, x
xa = np.hstack([[0],x,[0]])

# Get an array of ones and zeros, with ones for nonzeros of x and zeros elsewhere
xa1 =(xa!=0)+0

# Find consecutive differences on xa1
xadf = np.diff(xa1)

# Find start and stop+1 indices and thus the lengths of "islands" of non-zeros
starts = np.where(xadf==1)[0]
stops_p1 = np.where(xadf==-1)[0]
lens = stops_p1 - starts

# Mark indices where "minus ones" are to be put for applying cumsum
put_m1 = stops_p1[[stops_p1 < x.size]]

# Setup vector with ones for nonzero x's, "minus lens" at stops +1 & zeros elsewhere
vec = xa1[1:-1] # Note: this will change xa1, but it's okay as not needed anymore
vec[put_m1] = -lens[0:put_m1.size]

# Perform cumsum to get the desired output
out = vec.cumsum()

样例运行 -

In [116]: x
Out[116]: array([ 0. ,  2.3,  1.2,  4.1,  0. ,  0. ,  5.3,  0. ,  1.2,  3.1,  0. ])

In [117]: out
Out[117]: array([0, 1, 2, 3, 0, 0, 1, 0, 1, 2, 0], dtype=int32)

运行时测试 -

下面是一些运行时测试，将提议的方法与其他itertools.groupby基于的方法进行比较 -

In [21]: N = 1000000
    ...: x = np.random.rand(1,N)
    ...: x[x>0.5] = 0.0
    ...: x = x.ravel()
    ...: 

In [19]: %timeit sumrunlen_vectorized(x)
10 loops, best of 3: 19.9 ms per loop

In [20]: %timeit sumrunlen_loopy(x)
1 loops, best of 3: 2.86 s per loop

- Divakar

---

我也在这方面思考，但无法确定如何计算转换点处的负行程长度。感谢您提供的解决方案。可以通过删除vec1来缩短它，因为vec1.cumsum()等于(x != 0) + 0。我是新来的，不确定编辑别人的答案是否可以。 - steve

@steve 哦，非常好的观察！让我编辑一下，因为我还需要编辑解释步骤的文本。 - Divakar

再多一点... non_zero = (x != 0) + 0; xa = np.hstack([[0],non_zero,[0]]); xadf = np.diff(xa); 然后稍后 vec2 = non_zero - steve

@steve 在我进行了编辑之后，我也在考虑同样的事情，xadf 做了类似的操作。让我用这些优化再次编辑它。 - Divakar

@steve 现在这些编辑应该涵盖了那些优化。 - Divakar

@Divakar，有没有可能制作这个解决方案的反面？也就是只对零值进行累加求和？ - Amit Amola

Answer 2

你可以使用 itertools.groupby 和 np.hstack :

>>> import numpy as np
>>> x = np.array([2.3, 1.2, 4.1 , 0.0, 0.0, 5.3, 0, 1.2, 3.1])
>>> from itertools import groupby

>>> np.hstack([[i if j!=0 else j for i,j in enumerate(g,1)] for _,g in groupby(x,key=lambda x: x!=0)])
array([ 1.,  2.,  3.,  0.,  0.,  1.,  0.,  1.,  2.])

我们可以根据非零元素对数组元素进行分组，然后使用列表推导和枚举来替换非零子数组为其索引，最后使用np.hstack将列表展平。

- Mazdak

---

正如其他答案所提到的，由于它是纯Python而不是基于numpy的，因此对于大型输入来说，这个解决方案速度较慢。 - xjcl

Answer 3

这个子问题在我参加Kick Start 2021年A轮比赛时出现了。我的解决方案：

def current_run_len(a):
    a_ = np.hstack([0, a != 0, 0])  # first in starts and last in stops defined
    d = np.diff(a_)
    starts = np.where(d == 1)[0]
    stops = np.where(d == -1)[0]
    a_[stops + 1] = -(stops - starts)  # +1 for behind-last
    return a_[1:-1].cumsum()

事实上，这个问题也需要一个版本，其中您计算“向下”连续序列。因此，这里有另一个版本，带有可选的关键字参数rev=False，其结果与之前相同：

def current_run_len(a, rev=False):
    a_ = np.hstack([0, a != 0, 0])  # first in starts and last in stops defined
    d = np.diff(a_)
    starts = np.where(d == 1)[0]
    stops = np.where(d == -1)[0]
    if rev:
        a_[starts] = -(stops - starts)
        cs = -a_.cumsum()[:-2]
    else:
        a_[stops + 1] = -(stops - starts)  # +1 for behind-last
        cs = a_.cumsum()[1:-1]
    return cs

结果：

a = np.array([1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1])
print('a                             = ', a)
print('current_run_len(a)            = ', current_run_len(a))
print('current_run_len(a, rev=True)  = ', current_run_len(a, rev=True))

a                             =  [1 1 1 1 0 0 0 1 1 0 1 0 0 0 1]
current_run_len(a)            =  [1 2 3 4 0 0 0 1 2 0 1 0 0 0 1]
current_run_len(a, rev=True)  =  [4 3 2 1 0 0 0 2 1 0 1 0 0 0 1]

对于仅由0和1组成的数组，您可以将[0, a != 0, 0]简化为[0, a, 0]。但是，发布的版本也适用于任意非零数字。

- xjcl