如何从Pandas DataFrame中随机采样均衡的正负行对

问题背景 #

我有一个包含标签（label）、文件名（filepath）及其他元数据列的数据集，样本量最多可达200,000条。以下是模拟该数据集的代码片段：

import pandas as pd
import numpy as np
import IPython.display as ipd

size = 20000
df = []
rng = np.random.default_rng(0)
for i in range(size):
    l = rng.choice(('cat', 'dog', 'mouse', 'bird', 'horse', 'lion', 'rabbit'))
    fp = str(rng.integers(1e5)).zfill(6) + '.jpg'
    df.append((l, fp))
df = pd.DataFrame(df, columns=['label', 'filepath'])
ipd.display(df)

需求目标 #

希望高效生成N个随机数据对，要求数据集的正样本对与负样本对保持均衡。输出格式如下：

# df_out 的行数为 "N"
df_out = pd.DataFrame([], columns=['label_1', 'label_2', 'filepath_1', 'filepath_2'])

其中：

• 正样本对：label_1等于label_2的行对
• 负样本对：两个标签不相等的行对
  目标是让df_out中包含约50%的正样本对和50%的负样本对。

尝试的方法及问题 #

我尝试的第一种方法是从DataFrame中采样2N行，然后将它们合并为行对：

N = 20
ii = rng.permutation(np.arange(N*2)%len(df))
func = lambda x: x.dropna().astype(str).str.cat(sep=',')
df_out = df.iloc[ii].reset_index(drop=True)  # 子采样
df_out = df_out.groupby(df_out.index//2)  # 将每两行合并为一行
df_out = df_out.agg(func).reset_index(drop=True)  # 使用`func`合并行
for k in df.columns:
    df_out[[f'{k}_1',f'{k}_2']] = df_out[k].str.split(',', expand=True)
    del df_out[k]

这种方法可以生成行对，但无法保证正负样本对的均衡性：

# 正如预期，该比例不等于50%
print(sum(df_out.eval('label_1==label_2')) / N)

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