Pandas: Как вернуть один столбец из двух переданных в DataFrame.apply()?

Question

[Andrey M]

Добрый день!
Есть такой DataFrame:

Присутствуют пустые строки в столбце 'Возраст'. Я хочу их заменить на средние значения в соответствии с классом (т.к. есть корреляция).
Пробую реализовать так:

#Загрузка данных в data
data = ...

# Средний возраст для каждого соц класса
soc = pd.Series(data['Социальный класс'].unique(), name='Социальный класс')
mean = pd.Series([data['Возраст'].loc[data['Социальный класс'] == _].mean() for _ in soc], name='Возраст')
soc_mean = pd.concat([soc, mean], axis=1)

# Фильтруем дату. Если возраст NaN то выполняется условие получения возраста из soc_mean
func_filter = lambda x: soc_mean['Возраст'].loc[soc_mean['Социальный класс'] == x['Социальный класс']] if x['Возраст'] != x['Возраст'] else x['Возраст']
data['Возраст'] = data[['Социальный класс', 'Возраст']].apply(func_filter, axis=1)

Так как в последней строке я пытаюсь вернуть в data отфильтрованный soc_mean(DataFrame) в ячейку с NaN. Туда попадает весь DataFrame, вместо самого значения.
Однако второй столбец('Социальный класс') в soc_mean нужен для фильтрации, возможно ли вернуть только soc_mean['Возраст']?

PS lambda

Заметил, что лямбда очень нагружена и тяжела для понимания. Изначально думал над решение типа:
func = lambda x, y: soc_mean[x] if y != y else y
Но так и не додумался как это реализовать, потому что в .apply() не смог передать столбцы в разных переменных. Сколько не искал, так и не нашел примера где в .apply() передают разные столбцы по смыслу (1 для фильтра, 2 для изменения).

Answer 1

[Максим Припадчев]

Первое что бы получить mean по группам для этого в pandas есть мощный инструмент groupby
например в вашем случае
means = data.groupby('Социальный класс').mean()
Это даст вам датафрейм где вы на группе запустили аггрегирующуюю функцию (можно не одну)
теперь к ответу на вопрос все решается в одну строчку

data["Возраст"] = data.groupby("Социальный класс").transform(lambda x: x.fillna(x.mean()))

все этот код решит ваш вопрос, для его лучшего понимания (очень рекомендую) разберитесь как следует с groupby это очень частая проблема.

Comments

[Andrey M]

Спасибо большое, идею я понял, недавно только начал изучать pandas.
Единственное код выдавал ошибку, так как пытался записать несколько столбцов в 1.
Сработало так:

data['Возраст'] = data.groupby('Социальный класс')['Возраст'].transform(lambda x: x.fillna(x.mean()))

[Максим Припадчев]

Andrey M, Странно, при решении (хотя это классический паттерн) я создал игрушечный датафрейм и запустил код должно работать без ошибок и правильно. То есть для каждой группы nan заменяется на среднее значение этой группы. Честно я вот только что пересоздал датафрейм даже код опубликую

df = pd.DataFrame({
    'Social Class':np.random.randint(1,4,30),
    'Age': np.random.uniform(160,190,30)
})
positions_for_missing = np.random.randint(1,30,5)
df.loc[positions_for_missing,'Age'] = np.NaN
df["Age"] = df.groupby("Social Class").transform(lambda x: x.fillna(x.mean()))

Вот этот игрушечный код я использовал для ответа. По моему ваш похожий. И должен работать так же.

[Andrey M]

Максим Припадчев, Вы правы, ваш код работает. Но в моей таблице есть еще столбцы, я их убрал, т.к. смысловой нагрузки для вопроса они не имеют.
Так вот, стоит добавить еще 1 столбец и в Ваш тестовый DF, то в такой конструкции вылетает ошибка:
ValueError: Columns must be same length as key

[Максим Припадчев]

Andrey M, Да правильно. Мы то записываем в одну колонку, значит и запускать должны transform на одной колонке. Ну видите вы быстро учитесь, а мне следовало, сообразить сразу что если у вас еще колонки есть то вылезет ошибка и сделать ваш подправленный вариант. Так как он на мой код никакого влияния не окажет и он так же будет запускаться. Ну самое главное это конечно паттерн, потому что задача то ваша очень распространенная, запускать различные функции на группах. Так что на заметочку то этот паттерн возьмите.