如何按pandas DataFrame分组各列末尾值筛选列并生成嵌套列表

实现代码 #

import pandas as pd

# 示例数据
foo = pd.DataFrame({
                    'wag' : [1, 1, 1, 1, 44, 44, 44, 56, 56, 56, 56],
                     'A' :  [10, 8, 7, 3, 13, 22, 18, 43, 40, 38, 12],
                     'B' : [55, 48, 37, 34, 15, 12, 5, 1, 2, 2, 1]
                     })

foo_list = []
# 按wag字段分组遍历
for _, group in foo.groupby('wag', sort=False):
    # 提取除分组字段外的所有数值列
    value_cols = group.columns.drop('wag')
    # 取所有数值列的末尾值，找到值最小的对应列名
    col_min_last = group[value_cols].iloc[-1].idxmin()
    # 将选中列转为列表加入结果
    foo_list.append(group[col_min_last].tolist())

print(foo_list)
# 输出：[[10, 8, 7, 3], [15, 12, 5], [1, 2, 2, 1]]

说明 #

• 代码自动适配多列场景：所有非wag的列都会自动纳入对比，不需要手动指定列名
• 如果存在多个列的末尾值同为最小值，idxmin()会返回第一个出现的列，有自定义优先级需求可以调整这部分逻辑
• groupby加了sort=False参数，保证分组顺序和原数据中wag出现的顺序一致

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