Pandas - 从 groupby 操作的前 n 组创建新的 DataFrame

Question

具有以下DF：

   A   B
0  1  11
1  2  22
2  2  22
3  3  33
4  3  33

我想 groupby 'A' 然后取前 n 个组并从中创建一个新的数据框。 我环顾四周，找到了这个答案：

result = [g[1] for g in list(grouped)[:3]]

但是该解决方案返回一个列表而不是 DF，此外，从分组结果创建一个列表似乎是多余的。

更新：预期 output 是由前 n 个组组成的新 DF，例如，如果n=2 output 将是：

   A   B
0  1  11 <-- first group
1  2  22 <-- second group
2  2  22 <-- second group

任何帮助，将不胜感激

Answer 1

从技术上讲，您不能：这些组不一定按照您的 dataframe 的顺序排列： grouped结果按分组列排序（默认情况下，可以关闭），然后定义顺序。 换句话说，应使用分组列（此处为 A）中的值访问各个组。

在您的情况下，这可能有效：

import pandas as pd

df = pd.DataFrame({'A': [1, 2, 2, 3, 3], 'B': [11, 22, 22, 33, 33]})
grouped = df.groupby('A')
n = 2
df = pd.concat([group for name, group in grouped][:n])
print(df)

产生

   A   B
0  1  11
1  2  22
2  2  22

但如果输入 dataframe 如下（注意列中值的顺序）：

import pandas as pd

df = pd.DataFrame({'A': [2, 2, 3, 3, 1], 'B': [22, 22, 33, 33, 11]})
grouped = df.groupby('A')
n = 2
df = pd.concat([group for name, group in grouped][:n])
print(df)

前两个分组连接仍然是

   A   B
4  1  11
0  2  22
1  2  22

因为这些组是按“A”列中的值排序的。 （请注意这些值与以前的情况相同；但是，索引不同。）

因此，对于一组 groupby 结果，没有真正的“前 n 个元素”。

Answer 2

我们可以在这里使用pd.factorize和df.isin

ids = pd.factorize(df['B'])[1]
n = 2 # Take first two groups
m = df['B'].isin(ids[:n])
df.loc[m]

   A   B
0  1  11
1  2  22
2  2  22

当n=1时 Output

ids = pd.factorize(df['B'])[1]
n = 1 # Take first group
m = df['B'].isin(ids[:n])
df.loc[m]

   A   B
0  1  11

Answer 3

您可以获取索引并创建一个新的 dataframe ；

grouped = df.groupby('A')

假设 n = 2

indices = pd.Index.union(*[value 
                           for key, value in grouped.groups.items() 
                           if key in [*grouped.groups][:2]]
                         )

 indices
 Int64Index([0, 1, 2], dtype='int64')

 df.loc[indices]


   A   B
0  1  11
1  2  22
2  2  22

另请注意，如果您希望数据按特定顺序排列，则可以对分组进行排序； 如果sort is False ，它将根据出现在 dataframe 中的现有顺序返回 n 个组。