выбрать все элементы в группе, которые соответствуют определенному условию, используя панд

Question

У меня есть df, сгруппированный по id.Для каждой группы id я хотел бы вернуть все строки, в которых столбец a больше среднего значения a для этой группы.Я пробовал следующее:

df = pd.DataFrame(data = {"a": np.arange(10), "b": np.arange(10)[::-1], "c": np.random.choice(a = np.arange(10), size = 10)}, index = pd.Index(data = np.random.choice(a = [1,2,3], size = 10), name = "id"))
df.groupby("id").apply(lambda x: x[x.a > x.a.mean()])

Это вызывает ошибку ValueError: Дублированное имя уровня: "id", присвоенный уровню 1, уже используется для уровня 0.

Что я делаю неправильно

Answer 1

Используйте transform для того же Series как размер оригинала DataFrame для лучшей производительности, например apply решение:

df = df[df['a'] > df.groupby("id")['a'].transform('mean')]
print (df)
    a  b  c
id         
2   5  4  3
2   6  3  2
1   7  2  3
3   8  1  0
3   9  0  5

Деталь :

print (df.groupby("id")['a'].transform('mean'))
id
1    4.75
1    4.75
1    4.75
3    3.50
3    3.50
1    4.75
1    4.75
1    4.75
1    4.75
1    4.75
Name: a, dtype: float64

---

В вашем решении необходим параметр group_keys=False для избежания MultiIndex с такими же именами уровней, потому что id в имени индекса:

df = df.groupby("id", group_keys=False).apply(lambda x: x[x.a > x.a.mean()])

Если первыйreset_index() получить имя столбца id и имя индекса id, но есть те же значения:

df = df.reset_index().groupby("id").apply(lambda x: x[x.a > x.a.mean()])
print (df)
      id  a  b  c
id               
2  6   2  6  3  3
   7   2  7  2  9
   9   2  9  0  1
3  5   3  5  4  9
   8   3  8  1  8

Еще один тест - удаление index name - id:

df = df.rename_axis(None)
print (df)
   a  b  c
3  0  9  2
2  1  8  2
1  2  7  6
3  3  6  1
1  4  5  3
2  5  4  9
3  6  3  6
2  7  2  1
1  8  1  0
1  9  0  1

df = df.groupby(level=0).apply(lambda x: x[x.a > x.a.mean()])
print (df)
     a  b  c
1 1  8  1  0
  1  9  0  1
2 2  5  4  9
  2  7  2  1
3 3  6  3  6