如何用更Pythonic的方式实现Pandas按年统计同符号连续值达标次数？

我有一个由整数时间序列组成的pandas DataFrame，需要按年份对DataFrame分组，然后统计每年中同符号连续元素的绝对值之和≥5的次数。

示例数据与期望输出 #

>>> import pandas as pd
>>> l  = [1, -1, -4, 2, 2, 4, 5, 1, -3, -4]
>>> idx1 = pd.date_range('2019-01-01',periods=5)
>>> idx2 = pd.date_range('2020-01-01',periods=5)
>>> idx = idx1.union(idx2)
>>> df = pd.DataFrame(l, index=idx, columns=['a'])
>>> df
            a
2019-01-01  1
2019-01-02 -1
2019-01-03 -4  # 2019统计数=1：abs(-1)+abs(-4)≥5
2019-01-04  2
2019-01-05  2
2020-01-01  4
2020-01-02  5  # 2020统计数=1：abs(4)+abs(5)+abs(1)=10≥5
2020-01-03  1
2020-01-04 -3
2020-01-05 -4  # 2020统计数=2：abs(-3)+abs(-4)=7≥5

期望输出：

2019 1
2020 2

当前实现 #

我目前通过链式调用groupby和apply实现需求，编写了对应的分组函数和统计函数：

>>> import numpy as np
>>> def get_year(x):
        return x.year

>>> def count(group, t=5):
        c = 0 # 计数器
        s = 0 # 同符号连续元素的和                                                    
        for i in range(1,len(group)):
            if np.sign(group['a'].iloc[i-1]) == np.sign(group['a'].iloc[i]): 
                if s == 0:
                    s = group['a'].iloc[i-1] + group['a'].iloc[i]
                else:
                    s += group['a'].iloc[i]
                if i == (len(group) -1):
                   return c + 1
            elif (np.sign(group['a'].iloc[i-1]) != np.sign(group['a'].iloc[i])) and (abs(s) >= t):
                # 如果同符号连续序列中断且绝对值和≥t，计数器加1并重置s
                c += 1
                s = 0
            elif (np.sign(group['a'].iloc[i-1]) != np.sign(group['a'].iloc[i])) and (abs(s) < t):
                # 如果同符号连续序列中断且绝对值和<t，仅重置s
                s = 0
        return c

>>> by_year = df.groupby(get_year)
>>> by_year.apply(count)
2019 1
2020 2

我的问题是：是否存在更Pythonic的实现方式，能够得到相同结果且不依赖for循环？

可以利用pandas的矢量化操作和分组功能，完全避免显式for循环，代码更简洁高效：

步骤分解 #

• 标记同符号连续组：通过比较当前元素与前一个元素的符号，生成分组标签，每个连续同符号的序列属于同一个组。
• 计算每组的绝对值之和：按年份和同符号组分组，计算每组的绝对值总和。
• 统计符合条件的组数量：筛选出绝对值和≥5的组，再按年份统计数量。

完整代码 #

import pandas as pd
import numpy as np

# 生成示例数据
l  = [1, -1, -4, 2, 2, 4, 5, 1, -3, -4]
idx1 = pd.date_range('2019-01-01',periods=5)
idx2 = pd.date_range('2020-01-01',periods=5)
idx = idx1.union(idx2)
df = pd.DataFrame(l, index=idx, columns=['a'])

# 1. 添加年份列和符号列
df['year'] = df.index.year
df['sign'] = np.sign(df['a'])

# 2. 标记同符号连续组：当符号变化时，组号+1
df['group'] = (df['sign'] != df['sign'].shift()).cumsum()

# 3. 按年份和同符号组分组，计算每组绝对值之和
grouped_sum = df.groupby(['year', 'group'])['a'].apply(lambda x: x.abs().sum())

# 4. 筛选出和≥5的组，再按年份统计数量
result = grouped_sum[grouped_sum >=5].groupby('year').count()

print(result)

输出结果 #

year
2019    1
2020    2
dtype: int64

代码解释 #

• 标记同符号组：(df['sign'] != df['sign'].shift()).cumsum()会在符号变化时生成新的组号，确保每个连续同符号序列拥有唯一的组标识。
• 分组求和：按年份和组号双重分组，计算每组的绝对值总和，准确获取每个连续同符号序列的总和。
• 统计符合条件的组：筛选出总和≥5的组后，再按年份计数，得到每年的目标统计次数。

这种方法完全利用pandas的内置矢量化操作，避免了手动循环，代码更简洁易读，同时在大数据量下性能更优。

内容的提问来源于stack exchange，提问作者jgg