从 Pandas DataFrame 中有条件地抽取行

Question

我有一个pandas DataFrame，其中某些人被过度代表。我想进行子采样，并将每个人的观察数量限制在某个最大值以下。

目前我正在使用循环进行操作，并尝试从字典构建DataFrame。但是索引会成为问题，希望有人能指点一下更简单的解决方案。实际数据有约20K行，4K列和400人。谢谢。

示例数据。

df = pd.DataFrame({'name': ["Alice", "Alice", "Charles", "Charles", "Charles", "Kumar", "Kumar", "Kumar", "Kumar"],
              'height': [124, 125, 169, 178, 177, 172, 173, 175, 174]})

df
    height name
0   124 Alice
1   125 Alice
2   169 Charles
3   178 Charles
4   177 Charles
5   172 Kumar
6   173 Kumar
7   175 Kumar
8   174 Kumar

我现在的代码，尝试将每个人的行限制为2行。

sub_df = []
for name in pd.unique(df.name):
    sub_df.append(df[df.name == name].sample(n=2, random_state=42).to_dict())

pd.DataFrame(sub_df)

我正在了解的内容。

    height               name
0   {1: 125, 0: 124}    {1: 'Alice', 0: 'Alice'}
1   {2: 169, 3: 178}    {2: 'Charles', 3: 'Charles'}
2   {6: 174, 8: 175}    {6: 'Kumar', 8: 'Kumar'}

What I want.

    height name
0   125 Alice
1   124 Alice
2   169 Charles
3   178 Charles
4   174 Kumar
5   175 Kumar

- Chinntimes

Answer 1

在'name'上执行groupby，然后使用sample：

# groupby and sample
df = df.groupby('name').apply(lambda grp: grp.sample(n=2))

# formatting
df = df.reset_index(drop=True)

生成的输出结果：

   height     name
0     125    Alice
1     124    Alice
2     177  Charles
3     169  Charles
4     175    Kumar
5     173    Kumar

- root