Фильтр в группе по в pandas

#python #pandas #dataframe

#python #pandas #фрейм данных

Вопрос:

У меня есть следующий фрейм данных

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

И я хочу получить подмножество этого фрейма данных с такими группами g , что mean(x) > 0.6 . То есть я хочу filter_group , чтобы операция получила следующий фрейм данных:

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

 

Есть ли простой способ сделать это в pandas? Это похоже на having операцию в SQL, но немного отличается, поскольку я хочу получить фрейм данных с той же схемой, но меньшим количеством строк.


Для пользователей 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)
 

Ответ №1:

Используйте GroupBy.transform для повторного просмотра агрегированных значений по группам для возможного фильтрации исходных значений в boolean indexing :

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

Это решение лучше, если большой фрейм данных или много групп, если важна производительность:

 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)
 

Комментарии:

1. Интересный бенчмарк, удивительно, что наблюдается такое увеличение производительности (почти на 2 порядка !!!), Вы знаете, почему такая большая разница?

2. @DavidMasip — я думаю filter , что он разработан так, чтобы быть очень общим, поэтому медленным.

Ответ №2:

Посмотрев на это, альтернативный способ сделать это — использовать filter метод:

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

Для меня это имеет следующие преимущества:

  • Он легко обобщается, если в фильтре задействовано много столбцов.
  • Он использует парадигму цепных панд, которая мне нравится.