Python Pandas按公司行数匹配自定义权重生成新列方案咨询

针对你的需求，我们可以通过Pandas的分组（groupby）和自定义函数来实现，核心思路是按公司分组后，根据每组行数匹配权重字典，并确保最新日期对应最高权重。

步骤拆解 #

1. 数据预处理：确保日期列是datetime类型，方便后续排序。
2. 分组处理：按公司分组，对每个分组执行以下操作：
   • 将组内数据按日期降序排序（最新日期排在最前面）
   • 根据分组行数从权重字典中取出对应权重列表
   • 将权重列表按降序排序（保证最高权重对应最新日期），赋值给新列custom_weights
3. 合并结果：将处理后的所有分组合并为最终DataFrame

完整示例代码 #

import pandas as pd

# ----------------------
# 1. 准备示例数据（替换成你的真实数据）
# ----------------------
data = {
    'company': ['A']*7 + ['B']*4 + ['C']*2,
    'date': ['2024-01-01', '2024-01-02', '2024-01-03', '2024-01-04', '2024-01-05', '2024-01-06', '2024-01-07',
             '2024-01-01', '2024-01-02', '2024-01-03', '2024-01-04',
             '2024-01-01', '2024-01-02'],
    'value': [10,20,30,40,50,60,70, 100,200,300,400, 1000,2000]
}
df = pd.DataFrame(data)
df['date'] = pd.to_datetime(df['date'])  # 转换为日期类型

# ----------------------
# 2. 自定义权重字典（键=公司行数，值=对应权重列表）
# ----------------------
weight_dict = {
    7: [0.3, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2],
    4: [0.4, 0.3, 0.2, 0.1],
    2: [0.6, 0.4]
}

# ----------------------
# 3. 定义分组处理函数
# ----------------------
def assign_custom_weights(group, weight_map):
    # 组内按日期降序排序，最新日期优先
    sorted_group = group.sort_values('date', ascending=False).reset_index(drop=True)
    group_size = len(sorted_group)
    
    # 检查权重字典是否有对应行数的权重，避免KeyError
    if group_size not in weight_map:
        sorted_group['custom_weights'] = None  # 无对应权重时赋值None，也可改为抛出异常
        return sorted_group
    
    # 将权重按降序排列，确保最高权重匹配最新日期
    sorted_weights = sorted(weight_map[group_size], reverse=True)
    sorted_group['custom_weights'] = sorted_weights
    return sorted_group

# ----------------------
# 4. 应用到所有分组
# ----------------------
result_df = df.groupby('company', group_keys=False).apply(assign_custom_weights, weight_map=weight_dict)

# 查看结果
print(result_df[['company', 'date', 'custom_weights']])

关键细节说明 #

• 如果你的权重字典里的列表已经是按“最高到最低”顺序准备好的（比如weight_dict[7]就是给最新到最旧日期的权重），可以去掉sorted(..., reverse=True)，直接用weight_map[group_size]赋值，这样完全保留你的自定义权重顺序。
• 代码中加入了group_size not in weight_map的判断，避免出现KeyError，你可以根据需求改为抛出异常或者设置默认权重。
• group_keys=False确保分组后不会保留公司作为索引，结果更贴近原始DataFrame结构。

输出示例 #

company       date  custom_weights
0        A 2024-01-07             0.3
1        A 2024-01-06             0.2
2        A 2024-01-05             0.1
3        A 2024-01-04             0.1
4        A 2024-01-03             0.1
5        A 2024-01-02             0.1
6        A 2024-01-01             0.1
0        B 2024-01-04             0.4
1        B 2024-01-03             0.3
2        B 2024-01-02             0.2
3        B 2024-01-01             0.1
0        C 2024-01-02             0.6
1        C 2024-01-01             0.4

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