在DataFrame中获取N组中的M组

时间:2018-02-07 06:39:23

标签: python pandas dataframe group-by pandas-groupby

我有一个包含a列和b列的DataFrame。对于a列中的每个组,我希望获得列m的第一个n唯一值的b行。

举个简单的例子,假设mn分别为2和3。以下数据集:

          a  b
0     hello  1
1     hello  2
2     hello  0
3     hello  1
4     hello  3
5     hello  2
6     hello  3
7     hello  0
8   goodbye  0
9   goodbye  3
10  goodbye  1
11  goodbye  0
12  goodbye  2
13  goodbye  2
14  goodbye  1
15  goodbye  3

将转化为

          a  b
0     hello  1
1     hello  2
2     hello  0
3     hello  1
5     hello  2
7     hello  0
8   goodbye  0
9   goodbye  3
10  goodbye  1
11  goodbye  0
14  goodbye  1
15  goodbye  3

对于列a中的每个组,现在列m的前3个(n)唯一值的< = 2(b)个实例。

我希望找到比我想象的更优雅的groupby解决方案。我的解决方案使用两次调用groupby,然后屏蔽并连接结果, at best 很难阅读和理解。

In [1]: import pandas

In [2]: m = 2

In [3]: n = 5

In [4]: import random

In [5]: a = (['hello'] * 50) + (['goodbye'] * 50)

In [6]: b = list(range(10)) * 10

In [7]: random.shuffle(b)

In [8]: random.shuffle(a)

In [9]: df = pandas.DataFrame({'a': a, 'b': b})

# Ugly and un-elegant one-liner...
In [10]: pandas.concat(x[x.b.isin(x.b.unique()[:n])] for _, x in df.groupby(['a', 'b']).head(m).groupby('a')).sort_index()
Out[10]: 
          a  b
0     hello  5
1   goodbye  2
2     hello  8
3     hello  7
4     hello  4
5     hello  9
6     hello  7
7   goodbye  5
8   goodbye  8
9   goodbye  5
10  goodbye  0
11  goodbye  2
12  goodbye  3
16  goodbye  0
21    hello  9
22    hello  8
24  goodbye  3
27    hello  5
29  goodbye  8
46    hello  4

是否有一种更简洁的方法使用现有的DataFrame方法让我实现我想要的目标?

1 个答案:

答案 0 :(得分:3)

您可以创建两个条件,一个使用cumcount + groupby,另一个使用apply + isinunique + df) ,然后使用它们来索引m, n = 2, 3 i = df.groupby(['a', 'b']).cumcount().lt(m) j = df.groupby('a')['b'].apply(lambda x: x.isin(x.unique()[:n])) df[i & j] a b 0 hello 1 1 hello 2 2 hello 0 3 hello 1 5 hello 2 7 hello 0 8 goodbye 0 9 goodbye 3 10 goodbye 1 11 goodbye 0 14 goodbye 1 15 goodbye 3

createSL.Style = (Style)Application.Current.Resources["blueButton"];