如何将Pandas的groupby+pct_change(periods=30)语句转换为纯Numpy实现

原代码逻辑拆解 #

原Pandas语句df.groupby(['column'])['new_column'].pct_change(periods=30)的逻辑分为两步：

• 按column列对数据集做分组，不同分组的计算相互隔离
• 每个分组内对new_column列计算周期为30的百分比变化，计算公式为 当前值 / 向前偏移30位的数值 - 1，每组长度不足30的部分返回空值

纯Numpy实现代码 #

import numpy as np

def groupby_pct_change(group_col: np.ndarray, value_col: np.ndarray, periods: int = 30) -> np.ndarray:
    # 按分组列排序，同时保留原始索引用于最终还原顺序
    sorted_idx = np.lexsort((group_col,))
    sorted_groups = group_col[sorted_idx]
    sorted_values = value_col[sorted_idx]
    
    # 定位每个分组的起止边界，兼容数值、字符串类型的分组列
    prepend_val = sorted_groups[0] + 1 if np.issubdtype(sorted_groups.dtype, np.number) else object()
    append_val = sorted_groups[-1] + 1 if np.issubdtype(sorted_groups.dtype, np.number) else object()
    group_boundaries = np.diff(sorted_groups, prepend=prepend_val, append=append_val).nonzero()[0]
    
    # 初始化结果数组，默认值为nan
    result = np.full_like(value_col, np.nan, dtype=np.float64)
    
    # 逐分组计算pct_change
    for start, end in zip(group_boundaries[:-1], group_boundaries[1:]):
        group_length = end - start
        if group_length <= periods:
            continue
        # 计算百分比变化
        pct_vals = (sorted_values[start + periods: end] / sorted_values[start: end - periods]) - 1
        # 赋值到原始索引对应位置
        result[sorted_idx[start + periods: end]] = pct_vals
    
    return result

正确性验证 #

你可以用以下代码对比Numpy实现和原Pandas代码的输出：

import pandas as pd

# 构造测试数据
df = pd.DataFrame({
    "column": np.random.randint(0, 5, size=2000),
    "new_column": np.random.randn(2000).cumsum() + 10 # 加10避免出现0值导致inf
})

# 原Pandas输出
pd_result = df.groupby("column")["new_column"].pct_change(30).values
# Numpy实现输出
np_result = groupby_pct_change(df["column"].values, df["new_column"].values, periods=30)

# 验证非空位置数值完全一致
print(np.allclose(pd_result[~np.isnan(pd_result)], np_result[~np.isnan(np_result)]))
# 输出为True即说明实现正确

注意事项 #

• 如果分组列为字符串/类别类型，不需要额外预处理，上述代码已经做了兼容
• 如果存在大量极小分组，可将循环逻辑替换为向量化操作提升性能，常规业务场景下上述实现性能足够
• 若数值列存在0值，除以0会返回inf，和Pandas原生行为一致

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