在每組數據幀中獲取可變長度的行的最有效方法是什么

Question

為了清楚地說明我的問題，對於這樣的虛擬數據框：

df = pd.DataFrame({'X' : ['B', 'B', 'A', 'A', 'A'], 'Y' : [1, 2, 3, 4, 5]})

如何獲得組A的前1行和組B的前2行，並擺脫每個組的其余行？ 順便說一下，真實的數據集很大，有成千上萬的行和成千上萬的組。

輸出看起來像這樣：

pd.DataFrame({'X' : ['B', 'B', 'A'], 'Y' : [1, 2, 3]})

我的主要抱怨是.groupby().head()僅給我每個組中固定長度的行，並且我希望不同組中的行數不同。

Answer 1

一種方法是創建一個包含每個組應保留的行數的字典，在groupby.apply ，使用g.name作為鍵在字典中查找值，使用head方法可以保持每組不同的行：

rows_per_group = {"A": 1, "B": 2}    
df.groupby("X", group_keys=False).apply(lambda g: g.head(rows_per_group[g.name]))

#   X   Y
#2  A   3
#0  B   1
#1  B   2