Pandas按条件分组聚合列表：按用户会话筛选并保留Target聚合数组

原始数据 #

需求 #

1. 按id和session分组
2. 仅保留每个分组中target列至少包含2个1和2个0的分组
3. 符合条件的分组中，将每行的target值与smth列的数组组成子列表，再聚合为大列表（必须保留target用于还原原始数据）

实现代码 #

import pandas as pd
import numpy as np

# 构造原始DataFrame
data = {
    'id': [1,1,1,1,1,1,1,1,1,1,1,1,1,2],
    'target': [0,1,0,1,0,0,0,0,1,1,0,0,0,1],
    'session': [1,1,1,1,2,2,2,2,3,3,3,3,3,1],
    'smth': [np.array([1,2,3]), np.array([5,7,1]), np.array([2,3,4]), np.array([3,4,5]),
             np.array([2,2,8]), np.array([4,2,0]), np.array([0,0,0]), np.array([1,3,3]),
             np.array([1,4,4]), np.array([1,5,5]), np.array([1,6,6]), np.array([1,7,7]),
             np.array([1,8,3]), np.array([1,9,3])]
}
df = pd.DataFrame(data)

# 分组筛选+聚合
filtered_aggregated = (df.groupby(['id', 'session'])
                       # 筛选满足target有至少2个0和2个1的分组
                       .filter(lambda group: (group['target'].value_counts().get(0, 0) >= 2) 
                               and (group['target'].value_counts().get(1, 0) >= 2))
                       # 对符合条件的分组，将每行的target和smth转为子列表并聚合
                       .groupby(['id', 'session'])
                       .apply(lambda group: group[['target', 'smth']].values.tolist())
                       .tolist())

# 输出结果
for item in filtered_aggregated:
    print(item)

输出结果 #

[[0, array([1, 2, 3])], [1, array([5, 7, 1])], [0, array([2, 3, 4])], [1, array([3, 4, 5])]]
[[1, array([1, 4, 4])], [1, array([1, 5, 5])], [0, array([1, 6, 6])], [0, array([1, 7, 7])], [0, array([1, 8, 3])]]

代码说明 #

• filter方法：遍历每个分组，统计target中0和1的数量，只有两者都≥2的分组才会被保留
• apply方法：对筛选后的每个分组，提取target和smth列的数值，转为嵌套列表格式
• 最终通过tolist()将分组结果转为普通列表，得到符合需求的聚合结果

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