Pandas分组求和后如何获取含并列的最小N值项目名称

需求说明 #

现有模拟的DataFrame文本数据，需要按姓名分组求和各项目数值，找出每个用户对应求和值最小的2个项目名称（包含并列情况），期望输出：

David ['Project_C', 'Project_D']
Kate  ['Project_B', 'Project_D']
Mike  ['Project_A', 'Project_B', 'Project_D']

原代码及问题 #

原编写的Pandas代码如下：

import pandas as pd
from io import StringIO

csvfile = StringIO(
"""Name Project_A   Project_B   Project_C   Project_D
David   1       1   
David       1       
David   1   1       1
Mike    1       1   
Mike        1   1   
Kate    1       1   
Kate    1   1   1   1
Kate    1       1   
Kate        1       
""")

dict_p = {0 : 'Project_A',
1: 'Project_B',
2: 'Project_C',
3: 'Project_D'}

df = pd.read_csv(csvfile, sep = '\t', engine='python')

df = df.fillna(0)

df = df.groupby(['Name'])["Project_A", "Project_B", "Project_C", "Project_D"].apply(lambda x : x.astype(int).sum())

print (df)

# turn the interim result into a list, then index the smallest 2, and map from dict_p.
dict_from_df = df.set_index('Name').T.to_dict('list')
n_min_values = 2

for k,v in dict_from_df.items():
    index_of_smallest_two = sorted(range(len(v)), key=lambda k: v[k])[:n_min_values] 

    C = (pd.Series(index_of_smallest_two)).map(dict_p)
    final_list = list(C)

    print (k, final_list)

中间分组求和输出 #

Name  Project_A  Project_B  Project_C  Project_D
0  David          2          2          1          1
1   Kate          3          2          3          1
2   Mike          1          1          2          0

当前输出问题 #

David ['Project_C', 'Project_D']
Kate ['Project_D', 'Project_B']
Mike ['Project_D', 'Project_A']

Mike的Project_A和Project_B求和值均为1，但当前输出遗漏了Project_B——原代码仅截取排序后的前2个索引，未处理并列情况。

解决方案 #

核心思路：先计算每个用户的项目求和值，找到第N小的数值作为阈值，再筛选所有求和值≤该阈值的项目，自动包含所有并列项。

修正后的代码如下：

import pandas as pd
from io import StringIO

csvfile = StringIO(
"""Name Project_A   Project_B   Project_C   Project_D
David   1       1   
David       1       
David   1   1       1
Mike    1       1   
Mike        1   1   
Kate    1       1   
Kate    1   1   1   1
Kate    1       1   
Kate        1       
""")

df = pd.read_csv(csvfile, sep='\t', engine='python')
df = df.fillna(0)

# 分组求和，将Name设为索引简化后续操作
sum_df = df.groupby('Name')[["Project_A", "Project_B", "Project_C", "Project_D"]].sum().astype(int)

n_min = 2

for name, row in sum_df.iterrows():
    # 获取当前用户求和值的唯一值并排序，取第n_min小的数值作为阈值
    sorted_unique = sorted(row.unique())
    threshold = sorted_unique[n_min - 1] if len(sorted_unique) >= n_min else sorted_unique[-1]
    
    # 筛选所有求和值≤阈值的项目名称，可选排序让输出更规整
    min_projects = sorted(row[row <= threshold].index.tolist())
    
    print(f"{name:5} {min_projects}")

修正后输出 #

David ['Project_C', 'Project_D']
Kate  ['Project_B', 'Project_D']
Mike  ['Project_A', 'Project_B', 'Project_D']

代码改动说明 #

1. 分组求和时直接将Name设为索引，避免额外的索引转换操作，逻辑更简洁。
2. 对每个用户的求和行，提取唯一值并排序，精准定位第2小的数值作为筛选阈值（如Mike的求和值唯一排序后为[0,1,2]，阈值取1）。
3. 通过阈值筛选所有符合条件的项目，自动包含所有并列项（Mike的Project_D(0)、Project_A(1)、Project_B(1)均满足≤1的条件）。
4. 对项目名称排序，让输出格式更统一整洁。

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