如何按月份对Pandas DataFrame下采样，平衡TARGET=0与1的数量

原始数据情况 #

目标状态 #

需要让每个月份的TARGET=0和TARGET=1样本数完全相等，最终状态如下：

尝试的代码及问题 #

尝试了以下代码，但未删除任何数据，还触发了SettingWithCopyWarning警告：

for m in df['MONTH'].unique():
    number_of_ones = len(df[(df['MONTH']==m) & (df['TARGET']==1)])
    number_of_zeros = len(df[(df['MONTH']==m) & (df['TARGET']==0)])
    n_obs_to_drop = number_of_ones - number_of_zeros 
    df[df['MONTH']==m].drop(df[(df['MONTH']==m) & (df['TARGET']==1)].sample(n_obs_to_drop).index, inplace = True)

警告信息：

/opt/conda/envs/librerias_cbi/lib/python3.9/site-packages/pandas/core/frame.py:4901: SettingWithCopyWarning: 
A value is trying to be set on a copy of a slice from a DataFrame

See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy
  return super().drop(

补充可复现示例：

import pandas as pd
data = {
"MONTH": [202207, 202207, 202207, 202207, 202208, 202208, 202208, 202209, 202209, 202209, 202209],
"TARGET": [1, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0],
"other_column1": [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110],  # 示例额外列
"other_column2": [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100]
}
    
df = pd.DataFrame(data)

pd.crosstab(df['MONTH'],df['TARGET'])

运行后交叉表结果：

TARGET  0   1
MONTH       
202207  1   3
202208  1   2
202209  2   2

解决方案 #

方法1：修复原代码的问题 #

原代码无效的核心原因是df[df['MONTH']==m]返回的是原DataFrame的切片副本，对副本执行drop(inplace=True)不会修改原df。正确做法是收集所有需要删除的索引，再一次性从原df中删除：

drop_indices = []
for m in df['MONTH'].unique():
    # 获取当前月份TARGET=0和1的样本量
    zero_count = len(df[(df['MONTH'] == m) & (df['TARGET'] == 0)])
    one_count = len(df[(df['MONTH'] == m) & (df['TARGET'] == 1)])
    # 计算需要删除的TARGET=1样本数
    to_drop = one_count - zero_count
    if to_drop > 0:
        # 随机抽取要删除的TARGET=1样本索引
        drop_idx = df[(df['MONTH'] == m) & (df['TARGET'] == 1)].sample(n=to_drop).index
        drop_indices.extend(drop_idx)

# 从原df中删除目标索引
df = df.drop(drop_indices)

验证结果：

pd.crosstab(df['MONTH'], df['TARGET'])

输出：

TARGET  0  1
MONTH       
202207  1  1
202208  1  1
202209  2  2

方法2：分组采样实现（更简洁） #

利用groupby按月份分组，对每个组内的TARGET=1样本采样到和TARGET=0相同的数量，再合并所有组数据：

def downsample_group(group):
    zero_count = len(group[group['TARGET'] == 0])
    # 分离TARGET=0和TARGET=1的样本
    zeros = group[group['TARGET'] == 0]
    ones = group[group['TARGET'] == 1].sample(n=zero_count, random_state=42)  # random_state固定采样结果，可选
    # 合并返回平衡后的组数据
    return pd.concat([zeros, ones])

# 按MONTH分组并应用下采样函数
df_downsampled = df.groupby('MONTH', group_keys=False).apply(downsample_group)

这种方法无需手动处理索引，代码更简洁，同时保留所有其他列。

方法3：Pandas 1.3.0+ 版本的简化写法 #

如果你的Pandas版本≥1.3.0，可直接在groupby后使用支持lambda函数的sample方法：

df_downsampled = df.groupby(['MONTH', 'TARGET'], group_keys=False).apply(
    lambda x: x.sample(n=min(len(df[(df['MONTH'] == x.name[0]) & (df['TARGET'] == 0)]), len(x)))
)

这里x.name[0]取当前分组的月份，取该月TARGET=0的数量和当前组（TARGET=1）数量的较小值作为采样数，实现组内平衡。

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