如何高效获取DataFrame各列非零值的起止月份及中断重启记录

问题背景 #

给定如下结构的DataFrame（索引为月份）：

IndexMonth   Cus1   Cus2  Cus3  Cus4 ........ Cusn
2019-01       0     111    0      0           333
2019-02       0     111    0     666           0
2019-03      500     0     333    55           0  
2019-04      600     0     333    111          0
2019-05      600    100    0      111          0

需求：为每个客户列提取所有非零值的起始月份（包括中断后重新出现非零值的起始点）和终止月份，输出格式示例：

StartMonth          EndMonth
Cus1   2019-03             2019-05
Cus2   2019-01,2019-05     2019-02,2019-05
Cus3   2019-03             2019-04
Cus4   2019-02             2019-05
..
Cusn   2019-01             2019-01

高效实现方案 #

利用pandas的向量化操作和分组聚合，避免逐列逐行遍历，适合处理大量客户列的场景：

import pandas as pd

# 构造示例数据（实际使用时替换为你的DataFrame）
data = {
    'Cus1': [0, 0, 500, 600, 600],
    'Cus2': [111, 111, 0, 0, 100],
    'Cus3': [0, 0, 333, 333, 0],
    'Cus4': [0, 666, 55, 111, 111],
    'Cusn': [333, 0, 0, 0, 0]
}
df = pd.DataFrame(data, index=pd.date_range('2019-01', periods=5, freq='M').strftime('%Y-%m'))

# 1. 生成非零值标记矩阵
non_zero_mask = df != 0

# 2. 计算每个非零区间的分组ID：通过状态变化生成连续区间的唯一标识
group_ids = non_zero_mask.cumsum() - non_zero_mask.shift(fill_value=False).cumsum()
# 仅保留非零行的分组ID，零值行设为NaN
group_ids = group_ids.where(non_zero_mask)

# 3. 逐列处理，提取每个区间的起始/终止月份
result_list = []
for col in df.columns:
    col_groups = group_ids[col].dropna()
    if col_groups.empty:
        result_list.append({'Customer': col, 'StartMonth': '', 'EndMonth': ''})
        continue
    # 按分组ID聚合，取每个组的第一个和最后一个索引（即起始/终止月份）
    agg_result = col_groups.groupby(col_groups).agg(['first', 'last'])
    # 将多个区间合并为逗号分隔的字符串
    start_months = ','.join(agg_result['first'].index)
    end_months = ','.join(agg_result['last'].index)
    result_list.append({
        'Customer': col,
        'StartMonth': start_months,
        'EndMonth': end_months
    })

# 4. 转换为目标格式的DataFrame
final_result = pd.DataFrame(result_list).set_index('Customer')
print(final_result)

方案优势 #

• 核心逻辑用pandas内置的向量化操作（cumsum、shift）实现，比手动循环逐行判断快数倍，尤其适合列数多、数据量大的场景
• 分组聚合直接调用groupby.agg，避免手动遍历区间，代码简洁且性能稳定

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