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

Question

我正在尝试计算股票回报数据的连续上涨天数；因此，如果正日为 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 方式来实现我想要的，同时只浏览一次数据？我可以自己写一个循环，但想知道是否有更好的方法。

Answer 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

Answer 2

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

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

无论你做什么，Python 中的循环本质上都很慢。当然，如果您使用 pandas，那么您也在底层使用 numpy，并且具有所有性能优势。只是不要通过循环破坏它们。更不用说 Python 列表占用的内存比您想象的要多得多；可能比

8 bytes * length

要多得多，因为每个整数都可以包装到一个单独的对象中，并放置到内存中的一个单独区域中，并由列表中的指针指向。

numpy提供的向量化应该足够了，如果你能找到某种方法来表达这个函数而无需循环。事实上，我想知道是否有某种方法可以使用诸如A+B*C

之类的表达式来表示它。如果您可以用 

Lapack 中的函数构造此函数，那么您甚至有可能击败通过优化编译的普通 C++ 代码。

您还可以使用其中一种已编译的方法来加速循环。请参阅下面在 numpy 数组上使用

Numba 的解决方案。另一种选择是使用 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

事实上，上面第二个例子中的大部分时间都花在了 Just-In-Time 编译上。相反（请记住复制，因为函数会更改数组）。

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。

Answer 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

Answer 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]