数据帧内许多组的随机采样IF组大小大于一个值

Question

我是个新手。考虑以下内容。对于任何具有超过3个唯一SUB_ID的LOC_ID，我随机抽取3个唯一的SUB_ID样本，用concatenate:拼凑出一个新的数据帧

中的数据帧

df_concat = pd.DataFrame(columns=['LOC_ID', 'SUB_ID'])
for loc in test_df['LOC_ID'].unique():
sub_ids = test_df[test_df['LOC_ID'] == loc]['SUB_ID'].unique()
if len(sub_ids) > 3:
np.random.seed(622)
sub_ids_max = np.random.choice(sub_ids, size=3, replace=False)
df_sample = test_df[test_df['SUB_ID'].isin(sub_ids_max)]
else:
df_sample = test_df[test_df['SUB_ID'].isin(sub_ids)]
df_concat = pd.concat([df_concat, df_sample], ignore_index=True)

数据帧输出

以下是我还有一个6年的年份组成部分的真实情况。locids_sub中的每个LOC_ID至少在一年内具有10个以上唯一的sub_ID。根据LOC_ID，每年，我需要确保不超过10个唯一的SUB_ID:

max_subid = 10
years = df_bb['Year'].unique()
count_df = df_bb.groupby(['LOC_ID', 'Year']).nunique().reset_index()
locids_sub = count_df[count_df['SUB_ID'] > max_subid]['LOC_ID'].unique()
# Subset df_bb by locids_sub
df_bb_sub = df_bb[df_bb['LOC_ID'].isin(locids_sub)][['Year', 'LOC_ID', 'SUB_ID']]
df_concat = pd.DataFrame(columns=df_bb.columns)  # Initializing df
for year in years:
for loc in locids_sub:
sub_ids = df_bb_sub[(df_bb_sub['Year'] == year) & (df_bb_sub['LOC_ID'] == loc)]['SUB_ID'].unique()
if len(sub_ids) > max_subid:
np.random.seed(year+int(loc[-2:]))
sub_ids_max = np.random.choice(sub_ids, size=max_subid, replace=False)
df_sample = df_bb[df_bb['SUB_ID'].isin(sub_ids_max)]
else:
df_sample = df_bb[df_bb['SUB_ID'].isin(sub_ids)]
df_concat = pd.concat([df_concat, df_sample], ignore_index=True)

6年，locids_sub中有1460个LOC_ID，df_bb中有1828201行，这需要30分钟的运行时间
请告诉我如何提高效率。

Answer 1

您可以打乱整个未复制的数据帧并选择3行(相当快(：

(df.drop_duplicates(['LOC_ID', 'SUB_ID'])
.sample(frac=1).groupby('LOC_ID').head(n=3)

或groupby+apply(慢速(:

(df
.groupby('LOC_ID', group_keys=False)
.apply(lambda g: g if len(g2:=g.drop_duplicates('SUB_ID'))<3
else g2.sample(n=3))
)

输出：

LOC_ID  SUB_ID
5      a     903
4      a     918
3      a     888
8      b     314
7      b     721
6      b     231

注意。如果你喜欢这个，仍然可以考虑接受@Corralien的回答。如果没有他的回答，我可能不会尝试

Answer 2

IIUC，您可以使用：

out = (
df.drop_duplicates(['LOC_ID', 'SUB_ID'])
.assign(count=lambda x: x.groupby('LOC_ID').transform('nunique'))
.loc[lambda x: x.pop('count') >= 3].groupby('LOC_ID').sample(n=3)
)
print(out)
# Output
LOC_ID  SUB_ID
5      a     903
0      a     222
3      a     888
8      b     314
7      b     721
6      b     231

如果你有一个Year列，也许你可以把它改成以前的代码：

out = (
df.drop_duplicates(['Year', 'LOC_ID', 'SUB_ID'])
.assign(count=lambda x: x.groupby(['Year', 'LOC_ID']).transform('nunique'))
.loc[lambda x: x.pop('count') >= 3].groupby(['Year', 'LOC_ID']).sample(n=3)
)