DataFrame 的 Pandas 子集，每个类具有固定数量的值 [重复]

Posted 2023-02-16

【中文标题】DataFrame 的 Pandas 子集，每个类具有固定数量的值 [重复]

【英文标题】：Pandas subset of DataFrame with fixed number of value per class [duplicate]

【发布时间】：2019-10-08 16:19:31

【问题描述】：

我有一个 pandas DataFrame df，其中包含一个列 ID 和一个列 Type。这是一个例子：

print(df)

>>
+---------+---------+
|       ID|     Type|
+---------+---------+
|      AAA|        A|
|      BBB|        B|
|      CCC|        B|
|      DDD|        A|
|      EEE|        B|
|      FFF|        A|
|      GGG|        B|
+---------+---------+

从该 DataFrame 中，我想为每个 Type 提取一个带有 X 不同值的子 DataFrame。

这里是前面的例子（顺序无关紧要）：

X = 2
new_df = do_something(df, X)
print(new_df)

>>
+---------+---------+
|       ID|     Type|
+---------+---------+
|      AAA|        A|
|      DDD|        A|
|      BBB|        B|
|      CCC|        B|
+---------+---------+

有简单的方法吗？

【参考方案1】：

使用pandas.DataFrame.groupby.head:

import pandas as pd

df.groupby('Type').head(2)

输出：

    ID Type
0  AAA    A
1  BBB    B
2  CCC    B
3  DDD    A