在 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 范例。