如何在Python Pandas中为每个ID保留随机单月的数据行？

问题场景 #

给定如下DataFrame，需要为每个id随机保留某一个月份的所有行：

import pandas as pd

df = pd.DataFrame({'id': [1, 1, 1, 2, 2, 2],
                   'day': ['2023-03-01', '2023-03-10', '2023-04-05', '2023-03-15', '2023-04-20', '2023-04-25'],
                   'col1' : [123, 66, 7, 890, 456, 100]})

解决方案 #

可以通过提取年月标识、分组随机选择、筛选匹配行三个步骤实现，以下提供两种常用方法：

方法一：使用groupby+apply筛选

import numpy as np

# 1. 将日期转换为datetime格式，并提取年月（YYYY-MM）
df['year_month'] = pd.to_datetime(df['day']).dt.to_period('M')

# 2. 按id分组，为每个id随机挑选一个唯一的年月
selected_months = df.groupby('id')['year_month'].apply(lambda x: np.random.choice(x.unique()))

# 3. 筛选出每个id对应选中年月的所有行
result = df[df.apply(lambda row: row['year_month'] == selected_months[row['id']], axis=1)]

# 移除临时生成的year_month列（可选）
result = result.drop('year_month', axis=1)

print(result)

方法二：使用merge+query筛选（更直观）

import numpy as np

# 1. 提取年月标识
df['year_month'] = pd.to_datetime(df['day']).dt.to_period('M')

# 2. 生成每个id对应的随机年月，转为DataFrame
selected_df = df.groupby('id')['year_month'].unique()\
                .apply(np.random.choice)\
                .reset_index(name='selected_month')

# 3. 合并原始数据与选中年月，再筛选匹配行
result = df.merge(selected_df, on='id')\
           .query('year_month == selected_month')\
           .drop(['year_month', 'selected_month'], axis=1)

print(result)

代码解释 #

1. 提取年月标识：pd.to_datetime(df['day']).dt.to_period('M')将日期转为YYYY-MM格式的周期类型，确保同一月份的行被归为同一组。
2. 随机选择年月：通过groupby('id')对每个ID单独处理，用np.random.choice(x.unique())从该ID的所有不重复年月中随机选一个。
3. 筛选匹配行：两种方法分别通过apply逐行判断或merge+query的方式，保留每个ID对应选中年月的所有行。

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