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

什么是获得每个数据帧组的可变长度行的最有效方法

1 个答案: