如何在 Python 中的 pandas 数据帧上使用 groupby 创建多个列表聚合？

Question

以pandas DataFrame为例：

df = pd.DataFrame({"id": [1, 2, 3, 4, 5, 6],
                   "start": ["jan1", "jan1", "jan4", "feb17", "jan4", "mar3"],
                   "end": ["jan3", "jan3", "jan21", "feb17", "jan21", "mar4"],
                   "duration": [2, 2, 17, 0, 17, 1],
                   "case_id": ["case1", "case43", "case6", "case1", "case22", "case69"]
                  })

我想在 start、end 和 duration 列上使用 pandas groupby 操作来对数据框执行两个列表聚合：

• 每个组的id 值列表
• 每个组的case_id 值列表

我想要的输出如下所示：

start    end    duration    ids    cases
jan1     jan3   2           [1, 2] [case1, case43]
jan4     jan21  17          [3, 5] [case6, case22]
feb17    feb17  0           [4]    [case1]
mar3     mar4   1           [6]    [case69]

如何使用 pandas groupby 有效地做到这一点？

我知道如果我只需要一个聚合，我可以这样做：

df = df.groupby(['start', 'end', 'duration'])['id'].apply(list).to_frame()

如何为多个列表聚合执行此操作？如果有多种选择，什么是最耗时的？ （我正在转换的 DataFrame 非常大）

Answer 1

您需要使用pandas.groupby.agg，并将您要返回的列指定为list。

为了减少所需时间，因为您的数据中有分类列，请确保在 groupby 命令中使用 observed=True 选项。这确保它只创建存在条目的行（有关此here 的更多信息）

res = df.groupby(['start', 'end', 'duration'],observed=True)[['id','case_id']].agg(list).reset_index().sort_values(by='id')

输出：

res
Out[164]: 
   start    end  duration      id          case_id
1   jan1   jan3         2  [1, 2]  [case1, case43]
2   jan4  jan21        17  [3, 5]  [case6, case22]
0  feb17  feb17         0     [4]          [case1]
3   mar3   mar4         1     [6]         [case69]

假设您的唯一类别不是太多并且您的数据集不是太大，这应该不是问题。通常，处理字符串比处理数字需要更长的时间，因此如果运行时间过长，您可以尝试将您的 object 列转换为数字列并重新执行您的 groupby。