如何基于行条件对DataFrame分组应用自定义权重计算函数？

需求说明 #

按user字段对DataFrame分组，为每条记录分配权重，规则如下：

• 若用户出现3次及以上：
  • 首次（order=1）和末次（order=分组内总记录数）记录权重设为0.4
  • 其余记录权重为0.2/(分组内总记录数-2)
• 若用户出现不足3次，所有记录权重平均分配为1/分组内总记录数

预期结果 #

现有代码问题 #

现有代码已实现需求，但存在冗余的中间列（wt、grpsize）和重复的分组计算，希望优化：

• 改进groupby的使用方式，减少冗余步骤
• 让处理函数直接对分组内的元素进行计算，逻辑更紧凑

现有代码 #

import pandas as pd 
import numpy as np

# list of name, degree, score
user = ["josh", "josh", "josh", "liam", "liam", "liam", "abby", "abby"]
order = [1, 2, 3, 1, 3, 2, 1, 2]
device = ["a", "a", "a", "b", "a", "b", "c", "a"]
  
# dictionary of lists 
dict = {'user': user, 'order': order, 'device': device} 
    
df = pd.DataFrame(dict)
    

def func(T):
    if len(T) < 3:
        return 1/len(T)
    else:
        return 0.2/(len(T)-2)
        

df['wt']= df.groupby('user')['order'].transform(lambda x: func(x))
df['grpsize']= df.groupby('user')['order'].transform(lambda x: len(x))

df['final_wt'] = np.where(
    df['grpsize'] < 3, df['wt'], np.where(
    df['order'] ==1 , 0.4, np.where(
    df['order'] == df['grpsize'] , 0.4, df['wt'])))

优化后的实现 #

方案：使用groupby.apply直接处理分组 #

通过groupby.apply传入处理函数，直接在每个用户分组内完成权重计算，无需额外中间列，逻辑更清晰高效：

import pandas as pd

# 构造原始数据
user = ["josh", "josh", "josh", "liam", "liam", "liam", "abby", "abby"]
order = [1, 2, 3, 1, 3, 2, 1, 2]
device = ["a", "a", "a", "b", "a", "b", "c", "a"]
df = pd.DataFrame({'user': user, 'order': order, 'device': device})

def calculate_weights(group):
    group_size = len(group)
    if group_size < 3:
        # 不足3条记录，平均分配权重
        group['weight'] = 1 / group_size
    else:
        # 先给所有记录设中间值权重
        group['weight'] = 0.2 / (group_size - 2)
        # 覆盖首次和末次记录的权重
        group.loc[group['order'] == 1, 'weight'] = 0.4
        group.loc[group['order'] == group_size, 'weight'] = 0.4
    return group

# 应用分组计算，group_keys=False避免生成额外的分组索引
df = df.groupby('user', group_keys=False).apply(calculate_weights)
print(df)

优化点解析 #

1. 消除冗余步骤：不再生成wt、grpsize等中间列，直接在分组内完成所有权重计算
2. 直接处理分组元素：calculate_weights函数接收完整的用户分组，直接对分组内的每条记录进行权重赋值，逻辑连贯
3. 减少重复计算：原代码两次调用groupby.transform，优化后仅一次groupby.apply，降低了分组计算的开销
4. 可读性提升：权重规则的逻辑在函数内集中体现，更容易理解和维护

运行结果 #

user  order device  weight
0   josh      1      a     0.4
1   josh      2      a     0.2
2   josh      3      a     0.4
3   liam      1      b     0.4
4   liam      3      a     0.4
5   liam      2      b     0.2
6   abby      1      c     0.5
7   abby      2      a     0.5

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