Python中pandas DataFrame筛选后groupby保留原始索引问题

GroupBy.indices 属性返回的是分组对应当前DataFrame的内部行位置下标，而非保留的原始索引值，这是之前结果不符合预期的核心原因。

你可以直接对分组后的对象做索引聚合，再转为字典即可实现需求，示例代码如下：

import pandas as pd
import numpy as np

# 构造测试数据
headers = ['a','b']
nrows = 8
df = pd.DataFrame(columns = headers)
df['a'] = [0]*(nrows//2) + [1]*(nrows//2)
df['b'] = [2]*(nrows//4) + [4]*(nrows//4) + [2]*(nrows//4) + [4]*(nrows//4)

# 筛选目标子集
sub_df = df[df['a']==1]

# 提取分组对应原始索引
res = sub_df.groupby('b').apply(lambda group: group.index.to_numpy()).to_dict()
print(res)

输出结果为：{2: array([4, 5]), 4: array([6, 7])}，完全符合预期。

如果追求更高执行效率，也可以用agg替代apply实现相同效果：

res = sub_df.groupby('b').agg(index_list=('index', lambda x: x.to_numpy()))['index_list'].to_dict()

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