如何根據組標准使用Pandas對數據框進行子集化？

Question

我有一個大型數據集，具有以下結構

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函數執行此操作而不分組數據？ 我希望子集包含個別觀察。

Answer 1

作為@ 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 -length布爾系列用作選擇器：

>>> df.groupby("User")["X"].transform(sum) == 0
0     True
1     True
2    False
3    False
4    False
5     True
6     True
dtype: bool

當您想要通過縮減操作“廣播”組的結果備份到每個組的所有元素時，使用transform 。 它派上用場了。

Answer 2

DSM的答案是使用布爾掩碼選擇行，即使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