基于Pandas优化DataFrame中3→2状态转换及持续时长计算方案

问题说明 #

现有如下Pandas DataFrame，每行代表一个用户ID，每列对应不同月份的状态值：

• 初始状态默认为3
• 状态可变更为1、2，或变为nan（表示该ID被删除）

需求：针对每个ID，判断是否存在从状态3直接转换到2的情况，并计算转换后状态2的持续时长，最终期望输出结果为pd.Series([2, 3, 2, 1, 0, 0])。

原始数据 #

import pandas as pd
from numpy import nan

df = pd.DataFrame(data={
    "x1": [3, 3, 3, 3, 3, 3],
    "x2": [3, 3, 2, 2, 1, 3],
    "x3": [3, 2, 2, 3, 2, nan],
    "x4": [3, 2, 1, 2, 3, nan],
    "x5": [3, 2, 3, 2, 2, nan],
    "x6": [2, 1, nan, 2, 2, nan],
    "x7": [2, 2, nan, 2, 2, nan]
})

现有Numba实现方案 #

目前已有基于Numba的实现代码，执行结果符合预期：

import numba

@numba.njit
def _get_duration(l):
    counter = 0
    for i in range(1, len(l)):
        cond = l[i-1] == 3
        # 状态保持3，继续循环
        if cond and l[i] == 3:
            continue
        # 从3转换到2，初始化计数器
        elif cond and l[i] == 2:
            counter = 1
        # 状态保持2，计数器加1
        elif l[i-1] == 2 and l[i] == 2:
            counter +=1
        # 其他情况（转1、转nan或从非3/2状态变更），终止计算
        else:
            break
    return counter

@numba.njit
def get_stage2_duration(stg):
    N = stg.shape[0]
    return [_get_duration(stg[i]) for i in range(N)]

# 执行结果：[2, 3, 2, 1, 0, 0]
get_stage2_duration(df.values)

纯Pandas实现方案 #

以下是纯Pandas的实现，代码更简洁且效率优异：

高效向量化实现（推荐） #

完全避免逐行循环，利用Pandas的向量化操作处理，大数据量下性能更优：

# 1. 生成状态移位后的DataFrame（用于对比前后状态）
shifted = df.shift(axis=1)

# 2. 标记所有符合条件的连续2段：起始点是3→2的转换，后续持续为2
mask = df == 2
# 标记连续2段的起始位置
start_of_2 = mask & ~mask.shift(axis=1, fill_value=False)
# 筛选出起始点是3→2转换的有效段
valid_start = start_of_2 & (shifted == 3)

# 3. 为每个有效连续2段分配组ID
groups = mask.cumsum(axis=1)
valid_groups = groups.where(valid_start, 0).cummax(axis=1)

# 4. 统计每个有效组的长度，每行取第一个有效段的长度
group_counts = mask.groupby(valid_groups, axis=1).transform('count')
out = group_counts.max(axis=1).fillna(0).astype(int)

print(out)
# 输出：
# 0    2
# 1    3
# 2    2
# 3    1
# 4    0
# 5    0
# dtype: int64

逻辑说明 #

1. 先定位所有从3直接转换到2的位置，作为有效连续2段的起始点
2. 为每个连续的2段分配唯一组ID，仅保留有效起始点对应的组
3. 统计每个有效组的长度，每行取第一个有效段的长度即为结果；无有效转换的行结果为0

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