计算 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 方式来实现我想要的，同时只通过一次数据？我可以自己写一个循环，但想知道是否有更好的方法。

原文由 alex314159 发布，翻译遵循 CC BY-SA 4.0 许可协议

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

原文由 DSM 发布，翻译遵循 CC BY-SA 4.0 许可协议