python - 计算 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