计算 Python/pandas 数组中的连续正值

问题描述：

我正在尝试计算股票回报数据中连续上涨的天数；因此，如果正数为 1 而负数为 0，y=[0,0,1,1,1,0,0,1,0,1,1]则应该返回一个列表z=[0,0,1,2,3,0,0,1,0,1,2]。

我找到了一个只有几行代码但速度很慢的解决方案：

import pandas
y = pandas.Series([0,0,1,1,1,0,0,1,0,1,1])

def f(x):
    return reduce(lambda a,b:reduce((a+b)*b,x)

z = pandas.expanding_apply(y,f)

我猜我循环遍历整个列表的y次数太多了。有没有一种不错的 Pythonic 方法可以实现我想要的功能，同时只遍历数据一次？我可以自己写一个循环，但想知道是否有更好的方法。

解决方案 1：

>>> y = pandas.Series([0,0,1,1,1,0,0,1,0,1,1])

下面的内容可能看起来有点神奇，但实际上使用了一些常见的习语：由于pandas尚未对连续提供良好的本机支持groupby，因此您经常发现自己需要这样的东西。

>>> y * (y.groupby((y != y.shift()).cumsum()).cumcount() + 1)
0     0
1     0
2     1
3     2
4     3
5     0
6     0
7     1
8     0
9     1
10    2
dtype: int64

一些解释：首先，我们y与其自身的移位版本进行比较，以找出连续组何时开始：

>>> y != y.shift()
0      True
1     False
2      True
3     False
4     False
5      True
6     False
7      True
8      True
9      True
10    False
dtype: bool

然后（因为 False == 0 且 True == 1）我们可以应用累积和来获得各组的数字：

>>> (y != y.shift()).cumsum()
0     1
1     1
2     2
3     2
4     2
5     3
6     3
7     4
8     5
9     6
10    6
dtype: int32

我们可以使用groupby和cumcount来获得每个组中计数的整数：

>>> y.groupby((y != y.shift()).cumsum()).cumcount()
0     0
1     1
2     0
3     1
4     2
5     0
6     1
7     0
8     0
9     0
10    1
dtype: int64

添加一条：

>>> y.groupby((y != y.shift()).cumsum()).cumcount() + 1
0     1
1     2
2     1
3     2
4     3
5     1
6     2
7     1
8     1
9     1
10    2
dtype: int64

最后将起始值归零：

>>> y * (y.groupby((y != y.shift()).cumsum()).cumcount() + 1)
0     0
1     0
2     1
3     2
4     3
5     0
6     0
7     1
8     0
9     1
10    2
dtype: int64

解决方案 2：

如果某件事很明确，那么它就是“pythonic”的。坦率地说，我甚至无法使你的原始解决方案发挥作用。此外，如果它确实有效，我很好奇它是否比循环更快。你比较过吗？

现在，既然我们已经开始讨论效率，这里有一些见解。

不管你做什么，Python 中的循环本质上都很慢。当然，如果你使用 pandas，那么你也在底层使用 numpy，具有所有性能优势。只是不要通过循环破坏它们。更不用说 Python 列表占用的内存比你想象的要多得多；可能比 还多8 bytes * length，因为每个整数都可以包装成一个单独的对象并放置在内存中的单独区域中，并由列表中的指针指向。

如果您能找到某种方式来表达此函数而不使用循环， numpy提供的矢量化应该足够了。事实上，我想知道是否有某种方式可以使用诸如 之类的表达式来表示它。如果您可以使用LapackA+B*C中的函数构建此函数，那么您甚至可以击败使用优化编译的普通 C++ 代码。

您还可以使用其中一种编译方法来加速循环。请参阅下面使用Numba对 numpy 数组的解决方案。另一个选择是使用PyPy，尽管您可能无法将其与 pandas 正确结合使用。

In [140]: import pandas as pd
In [141]: import numpy as np
In [143]: a=np.random.randint(2,size=1000000)

# Try the simple approach
In [147]: def simple(L):
              for i in range(len(L)):
                  if L[i]==1:
                      L[i] += L[i-1]


In [148]: %time simple(L)
CPU times: user 255 ms, sys: 20.8 ms, total: 275 ms
Wall time: 248 ms


# Just-In-Time compilation
In[149]: from numba import jit
@jit          
def faster(z):
    prev=0
    for i in range(len(z)):
        cur=z[i]
        if cur==0:
             prev=0
        else:
             prev=prev+cur
             z[i]=prev

In [151]: %time faster(a)
CPU times: user 51.9 ms, sys: 1.12 ms, total: 53 ms
Wall time: 51.9 ms


In [159]: list(L)==list(a)
Out[159]: True

事实上，上面第二个例子中的大部分时间都花在了即时编译上。相反（记得复制，因为该函数会改变数组）。

b=a.copy()
In [38]: %time faster(b)
CPU times: user 55.1 ms, sys: 1.56 ms, total: 56.7 ms
Wall time: 56.3 ms

In [39]: %time faster(c)
CPU times: user 10.8 ms, sys: 42 µs, total: 10.9 ms
Wall time: 10.9 ms

因此，与简单版本相比，后续调用的速度提高了 25 倍。如果您想了解更多信息，我建议您阅读《高性能 Python》 。

解决方案 3：

与@DSM 的答案类似，但步骤更少：

s.groupby(s.ne(s.shift()).cumsum()).cumsum()

此表达式执行两个主要任务：

• s.ne(s.shift()).cumsum()或者(s != s.shift()).cumsum()为连续相同元素的运行分配唯一的组号

• s.groupby(...).cumsum()按这些数字分组s并计算每组内的累计和

输出：

0     0
1     0
2     1
3     2
4     3
5     0
6     0
7     1
8     0
9     1
10    2
dtype: int64

解决方案 4：

保持简单，使用一个数组、一个循环和一个条件。

a = [0,0,1,1,1,0,0,1,0,1,1]

for i in range(1, len(a)):
    if a[i] == 1:
        a[i] += a[i - 1]