У меня есть панд DataFrame df
, который содержит столбец ID
и столбец Type
. Вот пример:
print(df)
>>
+---------+---------+
| ID| Type|
+---------+---------+
| AAA| A|
| BBB| B|
| CCC| B|
| DDD| A|
| EEE| B|
| FFF| A|
| GGG| B|
+---------+---------+
Из этого DataFrame я хочу извлечь дополнительный DataFrame с X
различными значениями для каждого Type
.
Вот с предыдущим примером (порядок не имеет значения):
X = 2
new_df = do_something(df, X)
print(new_df)
>>
+---------+---------+
| ID| Type|
+---------+---------+
| AAA| A|
| DDD| A|
| BBB| B|
| CCC| B|
+---------+---------+
Есть ли простой способ сделать это?
Используйте 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