Pandas如何基于另一DataFrame的条件计算对应国家GDP均值

实现代码 #

你可以通过以下两种方式实现需求，第一种逻辑清晰易读，第二种为向量化实现，更适合大数据量场景：

方法1：遍历行计算（易理解）

import pandas as pd

# 构造df1和df2的代码你已经写好，这里直接复用即可
# 1. 将df1转为年份为索引、国家为列的宽表，存储对应年份各国GDP
gdp_wide = df1.pivot(index="year", columns="country", values="GDP")
country_list = ["France", "Germany", "US", "Canada"]

mean_gdp = []
for _, row in df2.iterrows():
    # 取当前年份的各国GDP，筛选出该组织当年的成员国GDP求平均
    current_year_gdp = gdp_wide.loc[row["year"]]
    member_gdp = current_year_gdp[row[country_list] == 1]
    mean_gdp.append(member_gdp.mean())

# 生成最终结果
df3 = df2.copy()
df3["mean_gdp"] = mean_gdp

方法2：向量化运算（效率更高）

import pandas as pd

gdp_wide = df1.pivot(index="year", columns="country", values="GDP")
country_list = ["France", "Germany", "US", "Canada"]

# 按年份对齐两个表的数据，计算成员国GDP总和后除以成员国数量得到平均值
merged = df2.merge(gdp_wide, on="year", suffixes=("_is_member", "_gdp"))
total_gdp = (merged[[c + "_is_member" for c in country_list]].values * merged[[c + "_gdp" for c in country_list]].values).sum(axis=1)
member_count = df2[country_list].sum(axis=1)

df3 = df2.assign(mean_gdp = total_gdp / member_count)

两种方法运行后得到的df3和你给出的预期输出完全一致。

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