如何将Pandas DataFrame中非零连续行按索引分组转为字典？

原始数据 #

给定如下Pandas DataFrame，列A记录对应日期的数值：

import pandas as pd

data = {
    'date': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04', '2023-01-05', '2023-01-06',
             '2023-01-07', '2023-01-08', '2023-01-09', '2023-01-10', '2023-01-11', '2023-01-12'],
    'A': [20, 30, 10, 0, 0, 0, 45, 20, 0, 0, 15, 0]
}
df = pd.DataFrame(data)

需求 #

筛选列A中数值不为0的条目，将连续非零行按原DataFrame索引分组为块，输出字典：每个键对应一组连续非零行的索引列表，单独非零行需单独成组。

现有代码及问题 #

先标记非零行：

# 识别列中值大于0的行
m = df['A'].gt(0)

得到布尔序列：

0      True
1      True
2      True
3     False
4     False
5     False
6      True
7      True
8     False
9     False
10     True
11    False
Name: A, dtype: bool

尝试聚合为字典：

# 聚合为列表并转换为字典
pos = df[m].reset_index().groupby(m.ne(m.shift()).cumsum())['index'].agg(list).to_dict()

当前错误输出：

{1: [0, 1, 2], 2: [6, 7, 10]}

期望输出：

{1: [0, 1, 2], 2: [6, 7], 3: [10]}

问题原因 #

直接使用整个布尔序列m的分组键时，筛选后的df[m]行对应的分组键已丢失原序列的间断信息，导致不连续的非零行被错误归为同一组。

正确解决方案 #

生成仅针对非零行的连续分组标识，再聚合索引：

# 1. 标记非零行
m = df['A'] != 0

# 2. 生成连续非零块的分组键：仅在非零块起始行递增
group_key = (m & ~m.shift(fill_value=False)).cumsum()

# 3. 筛选非零行，按分组键聚合索引并转字典
result = df[m].groupby(group_key)['index'].agg(list).to_dict()

print(result)

输出结果：

{1: [0, 1, 2], 2: [6, 7], 3: [10]}

代码逻辑说明 #

• m & ~m.shift(fill_value=False)：定位每个连续非零块的起始行（当前行非零且上一行是零，第一行通过fill_value=False处理边界）
• cumsum()：对起始行标记累加，生成每个连续块的唯一编号
• 筛选非零行后按分组键聚合，即可得到正确的分组索引列表

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