计算Python数组中的连续正值

时间:2014-12-23 19:09:59

标签: python pandas statistics

我试图计算股票回报数据的连续上升天数 - 因此,如果正数为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)

我猜测我在整个列表中循环次数太多次了。有没有一种很好的Pythonic方法可以实现我想要的,而只需要浏览一次数据?我自己可以写一个循环,但想知道是否有更好的方法。

谢谢!

4 个答案:

答案 0 :(得分:68)

这可能看起来有些神奇,但实际上使用了一些常见的习语:由于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

我们可以使用groupbycumcount为每个组中的整数计数:

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

答案 1 :(得分:5)

如果有什么东西是清楚的,那就是“pythonic”。坦率地说,我甚至无法使您的原始解决方案发挥作用。此外,如果它确实有效,我很好奇它是否比循环更快。你比较了吗?

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

无论你做什么,Python中的循环本质上都很慢。当然,如果你正在使用熊猫,你也会使用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

事实上,上面第二个例子中的大多数时间花在了即时编译上。相反(记得复制,因为函数会更改数组)。

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

因此,对于后续调用,与简单版本相比,我们的 25x-speedup 。如果您想了解更多信息,建议您阅读High Performance Python

答案 2 :(得分:1)

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

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]

答案 3 :(得分:-1)

为什么对超pythonic做事的痴迷?可读性+效率胜过“leet hackerz style。”

我会这样做:

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

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