如何基于分组计算为Pandas DataFrame添加城市人口占比列

我有如下结构的DataFrame：

我需要给这个DataFrame添加一列，显示每个城市人口占对应国家总人口的比例，具体要求：

• 添加名为relative population的第4列
• 先计算每个国家的总人口
• 新列的值为当前行城市人口除以对应国家的总人口

我不想遍历每一行，试了下面的代码：

# 计算指定国家的总人口
def TotalPop(thiscountry):
    return df[df['country'] == thiscountry]['population'].sum()

# 获取所有国家列表
countries_list = df['country'].unique()

for country in countries_list:
    df['population_Relative'] = df[df['country'] == country]['population'] / TotalPop(country)

但这段代码有问题：每次循环只会正确计算当前国家的population_Relative值，其他国家的对应行都会变成NaN。比如处理完Country 1时，它的城市相对人口是对的，但处理Country 2时，Country 1的对应值就变成NaN了。怎么才能不影响其他行实现需求？

正确实现方法：用groupby + transform #

Pandas的groupby配合transform可以高效完成这个需求，不需要循环，也不会覆盖其他行的数据：

# 按国家分组计算总人口，再广播到每一行，最后计算占比
df['relative population'] = df['population'] / df.groupby('country')['population'].transform('sum')

代码说明：

1. df.groupby('country')['population']：按country字段分组，提取每组的人口数据
2. .transform('sum')：对每组计算总人口，然后将该值复制到组内的每一行，得到一个和原DataFrame长度一致的Series
3. 用原population列除以这个Series，直接得到每个城市人口占对应国家总人口的比例

原代码问题分析 #

你每次循环时直接赋值df['population_Relative'] = ...，而右边的df[df['country'] == country]['population']仅包含当前国家的行，其他行都是NaN，所以每次赋值都会把之前计算好的其他国家数据覆盖成NaN。如果一定要用循环（不推荐，效率低），需要用.loc精准定位赋值：

# 先初始化列
df['population_Relative'] = 0.0
countries_list = df['country'].unique()

for country in countries_list:
    mask = df['country'] == country
    df.loc[mask, 'population_Relative'] = df.loc[mask, 'population'] / df.loc[mask, 'population'].sum()

但这种方法在数据量大时，性能远不如groupby+transform。

内容的提问来源于stack exchange，提问作者terauser