Pandas按ID分组并拼接vector列列表元素的技术需求

原始数据 #

我们有如下Pandas DataFrame：

需求 #

按ID分组，将每个ID对应的所有vector列表拼接成一个完整列表，得到如下结果：

解决方案 #

使用Pandas的groupby结合agg方法，通过自定义逻辑实现列表拼接：

import pandas as pd

# 构造原始DataFrame
data = {
    'ID': [0, 0, 1, 1, 1, 2],
    'time': [81219, 32323, 423, 4232, 42424, 42422],
    'vector': [['N1', 'N2', 'N3'], ['N1', 'N4'], ['N2', 'N3'], ['N1'], ['N2', 'N2'], ['N3', 'N5']],
    'other': ['blue', 'green', 'red', 'blue', 'red', 'blue']
}
df = pd.DataFrame(data)

# 分组拼接vector列表
result_df = df.groupby('ID')['vector'].agg(lambda x: [item for sublist in x for item in sublist]).reset_index()

print(result_df)

代码说明 #

1. groupby('ID')['vector']：按ID字段对vector列分组
2. agg(lambda x: [item for sublist in x for item in sublist])：用列表推导式将分组内的所有子列表展开并拼接成一个大列表
3. reset_index()：将分组后的索引转为普通列，得到结构规整的结果DataFrame

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