在 pandas 中按分組過濾

Question

我有以下 dataframe

 df = pd.DataFrame(dict(g = [0, 0, 1, 1, 2, 2], x = [0, 1, 1, 2, 2, 3]))

我想獲得這個 dataframe 的一個子集，其中的組來自g使得mean(x) > 0.6 。 也就是我想要一個filter_group操作得到如下dataframe：

>>> filtered_df = filter_group(df)
>>> filtered_df
   g  x
2  1  1
3  1  2
4  2  2
5  2  3

在 pandas 中是否有一種簡單的方法可以做到這一點？ 這類似於 SQL having操作，但有點不同，因為我想獲得具有相同架構但行數更少的 dataframe。

---

對於 R 用戶，我想做的是：

library(dplyr)
df <- tibble(
  g = c(0, 0, 1, 1, 2, 2),
  x = c(0, 1, 1, 2, 2, 3)
)

df %>% 
  group_by(g) %>% 
  filter(mean(x) > 0.6)

Answer 1

使用GroupBy.transform為每組重復聚合值，以獲得boolean indexing中可能的過濾器原始值：

df[df.groupby('g')['x'].transform('mean') > 0.6]

如果大型 DataFrame 或許多組如果性能很重要，則此解決方案更好：

np.random.seed(2020)

N = 10000
df = pd.DataFrame(dict(g = np.random.randint(1000, size=N), 
                       x = np.random.randint(10000, size=N)))
print (df)
        

In [89]: %timeit df[df.groupby('g')['x'].transform('mean') > 0.6]
2.01 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [90]: %timeit df.groupby('g').filter(lambda df: df['x'].mean() > 0.6)
145 ms ± 2.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

通過查看它，另一種方法是使用filter方法：

df.groupby('g').filter(lambda df: df['x'].mean() > 0.6)

對我來說，這有以下優點：

• 如果過濾器中涉及許多列，它很容易概括。
• 它使用我喜歡的鏈式 pandas 范例。