Заполнение столбца по условию в другом столбце

Есть DataFrame:

дф

Хотелось бы заполнить пропуски в столбце 'living_area' медианой, учитывая количество комнат в квартире (столбец 'rooms').

Попытался реалиизовать функцию, но нормально не заработала, так как исправляла весь df, а не конкретный столбец. Попробовал совсем топорным способом заполнить, чтобы понять процесс и вроде получилось:

stud = data[data['rooms']==0]
stud #192 студии
stud['living_area'].median() #медиана студий - 18 метров
stud['living_area'] =  stud['living_area'].fillna(stud['living_area'].median())
stud.isnull().sum() #192 пропуска заполенены медианой 
data[data['rooms']==0] = stud
data.isnull().sum()
data[data['rooms']==0].isnull().sum()

Как такую операцию можно красиво записать в функцию для определенного столбца через apply с циклом из значений всех квартир? То есть, в таком цикле:

for lsquare in list(data['rooms'].unique()):

Ответы (2 шт):

Автор решения: MaxU

Воспользуйтесь sklearn.impute.SimpleImputer

Пример:

from sklearn.impute import SimpleImputer

imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
data.loc[data['rooms'] == 0, 'living_area'] = \
    imp_mean.fit_transform(data.loc[data['rooms'] == 0, 'living_area'])
→ Ссылка
Автор решения: Andrey Ilin

Сначала нужно получить медианное значение 'living_area' в зависимости от комнат. Затем итерируемся по кол-ву комнат и подставляем в пропущенные значения нужную медиану.

 l_room_median = train.pivot_table(values='living_area', index='rooms', aggfunc='median')
 for y in (l_room_median.index):
    train.living_area= np.where((pd.isnull(data['living_area'])) & \
    (train['rooms'] == y), int(l_room_median.loc[y]), int(train['living_area'])
→ Ссылка