如何根据组标准使用Pandas对数据框进行子集化?

时间:2015-01-09 19:44:09

标签: python pandas

我有一个包含以下结构的大型数据集

User     X
1        0
1        0
2        0
2        0
2        1
3        0
3        0

我想获取一部分数据,使每个用户的X列总和为0.鉴于上述示例,该子集应仅包括用户1和3的观察结果,如下所示

User     X
1        0
1        0
3        0
3        0

有没有办法在不分组数据的情况下使用groupby函数执行此操作?我希望子集包含个别观察。

2 个答案:

答案 0 :(得分:8)

作为@ unutbu答案的替代方案,还有

>>> df.loc[df.groupby("User")["X"].transform(sum) == 0]
   User  X
0     1  0
1     1  0
5     3  0
6     3  0

这将创建一个df - 长度布尔系列用作选择器:

>>> df.groupby("User")["X"].transform(sum) == 0
0     True
1     True
2    False
3    False
4    False
5     True
6     True
dtype: bool
当您想要将groupby减少操作的结果“广播”回到每个组的所有元素时,使用

transform。它派上用场了。

答案 1 :(得分:5)

使用布尔掩码选择行的

DSM's answer即使DataFrame具有非唯一索引也能正常工作。 我的方法使用索引值选择行,当索引是唯一的时稍慢,而当索引包含重复值时明显更慢

@roland:请考虑接受DSM的回答。


您可以使用groupby-filter

In [16]: df.loc[df.groupby('User')['X'].filter(lambda x: x.sum() == 0).index]
Out[16]: 
   User  X
0     1  0
1     1  0
5     3  0
6     3  0

groupby-filter本身只返回:

In [29]: df.groupby('User')['X'].filter(lambda x: x.sum() == 0)
Out[29]: 
0    0
1    0
5    0
6    0
Name: X, dtype: int64

但您可以使用其索引

In [30]: df.groupby('User')['X'].filter(lambda x: x.sum() == 0).index
Out[30]: Int64Index([0, 1, 5, 6], dtype='int64')

使用df.loc选择所需的行。


以下是我使用的基准:

In [49]: df2 = pd.concat([df]*10000)   # df2 has a non-unique index

Ctrl - C '这是因为它花了太长时间才完成:

In [50]: %timeit df2.loc[df2.groupby('User')['X'].filter(lambda x: x.sum() == 0).index]

当我意识到自己的错误时,我创建了一个具有唯一索引的DataFrame:

In [51]: df3 = df2.reset_index()     # this gives df3 a unique index

In [52]: %timeit df3.loc[df3.groupby('User')['X'].filter(lambda x: x.sum() == 0).index]
100 loops, best of 3: 13 ms per loop

In [53]: %timeit df3.loc[df3.groupby("User")["X"].transform(sum) == 0]
100 loops, best of 3: 11.4 ms per loop

这表明即使使用非唯一索引,DSM的方法也能很好地运行:

In [54]: %timeit df2.loc[df2.groupby("User")["X"].transform(sum) == 0]
100 loops, best of 3: 11.2 ms per loop