计算 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，尽管您可能无法将它与熊猫正确结合。

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 倍的加速。如果您想了解更多信息，我建议您阅读High Performance Python。

Answer 3

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

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]

Answer 4

为什么对超蟒蛇式的做事方式如此着迷？可读性 + 效率胜过“leet hackerz 风格”。

我会这样做：

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