计算Python数组中的连续正值

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)

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

谢谢!

Answer 1

这可能看起来有点神奇,但实际上使用了一些常见的习语:因为pandas对于连续的还没有很好的原生支持groupby,你经常会发现自己需要这样的东西.

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

一些解释:首先,我们比较y自身的移位版本,以找出连续组开始的时间:

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

然后(因为False == 0和True == 1)我们可以应用累积和来获得组的数字:

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

我们可以使用groupby和cumcount得到每个组中的整数计数:

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

添加一个:

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

Answer 2

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

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

无论你做什么,Python中的循环本质上都很慢.当然,如果你正在使用熊猫,你也会使用numpy,具有所有的性能优势.只是不要通过循环来销毁它们.这并不是说Python列表比你想象的要多得多; 可能远远超过8 bytes * length,因为每个整数可以被包装到一个单独的对象中并放入内存中的一个单独区域,并由列表中的指针指向.

numpy提供的矢量化应该足够了,如果你能找到一些方法来表达这个功能而不需要循环.事实上,我想知道是否有某种方式通过使用表达式来表示它A+B*C.如果你可以用Lapack中的函数构造这个函数,那么你甚至可以击败用优化编译的普通C++代码.

您还可以使用其中一种编译方法来加速循环.请参阅下面的numby数组中的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

因此,对于后续调用,与简单版本相比,我们的速度提高了25倍.如果您想了解更多,我建议您阅读高性能Python.