Замените мутате (пакет dplyr) в python pandas

Question

Есть ли функция, аналогичная mutate (dplyr), с помощью которой я могу добавить новый столбец сгруппированных данных, применив функцию на одном из столбцов сгруппированных данных? Ниже приводится подробное описание проблемы:

Я генерируемые выборки данных с помощью следующего кода

x<- data.frame(country = rep(c("US", "UK"), 5), state = c(letters[1:10]), pop=sample(10000:50000,10)) 

Теперь я хочу, чтобы добавить новый столбец, который имеет максимальное население для США и Великобритании. Я могу сделать это с помощью следующей функции в R

x<- group_by(x, country) 
x<- mutate(x,max_pop = max(pop)) 
x<- arrange(x, country) 

Так что мой вопрос, как я могу сделать это в Python с помощью панд. Я пробовал следовать, но это не сработало

x['max_pop'] = x.groupby('country').pop.apply(max) 

Answer 1

Вы хотите использовать transform. transform вернет объект с тем же индексом, что и сгруппированный, что позволяет легко назначить обратно в качестве нового столбца в этом объекте, если он является фреймворком данных.

x['max_pop'] = x.groupby('country').pop.transform('max') 

---

Настройка

import pandas as pd 

x = pd.DataFrame(dict(
    country=['US','UK','US','UK'], 
    state=['a','b','c','d'], 
    pop=[37088, 46987, 17116, 20484] 
))