如何在Pandas DataFrame中计算滚动30天内国家的稀有度得分

我有一个包含country列和datetime列的Pandas DataFrame，需要为每行计算其对应国家在滚动30天窗口内的稀有度得分，规则如下：

1. 先计算该国在过去30天数据中的出现占比
2. 稀有度得分是所有与该行国家稀有度相同或更稀有的国家的出现占比总和（稀有度越高=占比越低）

示例数据 #

import pandas as pd

country_data = {
    'country': ['USA', 'USA', 'USA', 'Canada', 'UK', 'Canada', 'USA', 'Canada', 'Canada', 'UK'],
    'datetime': [
        '2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04', '2023-01-05',
        '2023-01-06', '2023-01-10', '2023-01-11', '2023-01-12', '2023-02-07'
    ],
}
country_df = pd.DataFrame(country_data)

期望输出 #

country datetime    rarity_score
USA     2023-01-01  1.0 
USA     2023-01-02  1.0
USA     2023-01-03  1.0
Canada  2023-01-04  0.25
UK      2023-01-05  0.4
Canada  2023-01-06  0.5
USA     2023-01-10  1.0
Canada  2023-01-11  0.5
Canada  2023-01-12  1.0
UK      2023-02-07  0.5

得分解释 #

• 第1-3行：USA在过去30天内是唯一出现的国家，占比100%，得分1.0
• 第4行：过去30天共4条数据，Canada出现1次占25%，无更稀有国家，得分0.25
• 第5行：过去30天共5条数据，UK与Canada各出现1次，占比均为20%，总和为0.4
• 第6行：过去30天内Canada出现2次占40%，UK出现1次占10%，两者占比总和为0.5
• 第7行：USA在过去30天出现4次，其他国家占比均更低，总占比100%，得分1.0
• 第8行：Canada在过去30天出现3次，UK出现1次，两者占比总和为0.5
• 第9行：Canada和USA在过去30天均出现4次，UK出现1次，总占比100%，得分1.0
• 第10行：2023-02-07往前30天内仅包含前4行数据，UK与USA各出现1次，占比总和50%，得分0.5

步骤1：预处理日期列 #

先将datetime列转换为Pandas日期类型，确保窗口计算的准确性：

country_df['datetime'] = pd.to_datetime(country_df['datetime'])

步骤2：自定义滚动窗口稀有度计算函数 #

通过apply遍历每行，计算对应30天窗口内的稀有度得分：

def calculate_rarity(row, df):
    # 确定30天窗口的时间范围
    window_end = row['datetime']
    window_start = window_end - pd.Timedelta(days=30)
    # 筛选窗口内的所有数据
    window_data = df[(df['datetime'] >= window_start) & (df['datetime'] <= window_end)]
    
    if len(window_data) == 0:
        return 0.0
    
    # 计算窗口内各国的出现占比，并按占比升序排序
    country_ratios = window_data['country'].value_counts(normalize=True).sort_values()
    # 获取当前国家的占比
    current_ratio = country_ratios.get(row['country'], 0.0)
    # 求和所有占比≤当前占比的国家的占比总和（占比越低越稀有）
    return country_ratios[country_ratios <= current_ratio].sum()

# 应用函数生成稀有度得分列
country_df['rarity_score'] = country_df.apply(calculate_rarity, df=country_df, axis=1)

步骤3：查看结果 #

运行以下代码即可得到与期望一致的输出：

print(country_df[['country', 'datetime', 'rarity_score']].to_string(index=False))

优化提示 #

如果处理大规模数据，apply的效率可能不足，可以考虑用滚动窗口分组结合广播的方式优化；中小规模数据下，上述方法简洁直观，足以满足需求。

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