在 for 循环中保存和删除数据帧

Question

我有许多数据帧，我必须从中抽取样本。从该数据帧中获取的样本必须从下一个数据帧中排除，以免出现任何“双”样本，因为存在一些重叠。

我的代码如下

df_list = [df1, df2, df3, df4, df5]
samplesizes = [8, 2, 4, 4, 2]

sample = []

for df, samplesize in zip(df_list, samplesizes):
    if sample: #can't drop in the first loop
        df = df.drop(sample) #I want to drop the taken samples from the current df

    if max_pop_size < len(df):
        samplesize = max_pop_size #can't take a sample larger than population

    sample.append(df.sample(samplesize, random_state=1000))

我在第一个循环后陷入了下降。我已经尝试了几种方法，但似乎都不起作用。

任何帮助将不胜感激！

Answer 1

这是一个有趣但并不简单的问题。

我认为你不能/不应该：

• 预过滤重复项，因为它可能会删除太多行（它可能会删除第 2 组中第 1 组中的行，即使它不是从第 1 组中采样的）
• 采样然后进行后过滤，因为采样后删除行会以不可预测的方式减少样本大小
• 在循环中连接采样的 DataFrame，这具有二次复杂度，并且会使代码对于大量输入变慢

我认为合理的方法可能是首先计算 ID 来识别重复的行，然后循环遍历组并跟踪已采样的行：

df1 = pd.DataFrame({'A': [1,2,3], 'B': [10,20,30]})
df2 = pd.DataFrame({'A': [2,3,4], 'B': [20,30,40]})
df3 = pd.DataFrame({'A': [1,4,5], 'B': [10,40,50]})
np.random.seed(1)

df_list = [df1, df2, df3]
samplesizes = [2, 1, 1]

# combine all datasets to be able to identify the duplicates
tmp = (pd.concat(df_list, keys=range(len(df_list)))
         .assign(ID=lambda d: d.groupby(list(tmp)).ngroup())
         .set_index('ID', append=True)
      )

seen = set()

sample = []
for k, g in tmp.groupby(level=0, sort=False):
    selected = g.drop(index=seen, level='ID').sample(samplesizes[k])
    seen.update(selected.index.get_level_values('ID'))
    sample.append(selected.droplevel([0, -1]))

print(sample)

当然，在某些情况下，可能没有足够的行可供采样，但如果您想保持无偏采样，这并不是真正可以避免的。

输出示例：

[   A   B
 0  1  10
 2  3  30,
    A   B
 2  4  40,
    A   B
 2  5  50]