Pandas中DataFrame列间除法报错KeyError，求原因及解决

需求 #

有两个结构相似的DataFrame，要把df1的count列除以df2的count列，在df1里新增一列叫weights。

示例数据 #

dict1 = {'category': {0: 0.0, 1: 1.0, 2: 0.0, 3: 0.0, 4: 1.0}, 'Id': {0: 24108, 1: 24307, 2: 24307, 3: 24411, 4: 24411}, 'count': {0: 3, 1: 2, 2: 33, 3: 98, 4: 33}}
df1 = pd.DataFrame(dict1)
dict2 = {'Id': {0: 24108, 1: 24307, 2: 24411}, 'count': {0: 3, 1: 35, 2: 131}}
df2 = pd.DataFrame(dict2)

你的代码和报错 #

你写的代码：

df1['weights'] = (df1['count'].div(df1.merge(df2, on = 'Id', how = 'left')['count'].to_numpy()))

报错信息：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/opt/conda/lib/python3.8/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3360             try:
-> 3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:

/opt/conda/lib/python3.8/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

/opt/conda/lib/python3.8/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'count'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
/tmp/ipykernel_354/1318629977.py in <module>
      1 complete['weights'] = (complete['count']
----> 2                       .div(complete.merge(totals, on = 'companyId', how = 'left')['count'].to_numpy())
      3                       )

错误原因 #

问题出在合并操作上：当两个DataFrame有同名列（这里都是count），pandas会自动给合并后的列加后缀区分——来自df1的是count_x，来自df2的是count_y，原来的count列名已经不存在了，你直接访问['count']自然会触发KeyError。

另外注意报错里的变量名是complete和totals，和你示例里的df1、df2不一样，但问题本质完全相同。

解决办法 #

办法1：合并时指定后缀，明确引用列 #

# 合并时指定后缀，清楚区分两个count列
merged_df = df1.merge(df2, on='Id', how='left', suffixes=('_df1', '_df2'))
# 用df1的count除以df2的count
df1['weights'] = merged_df['count_df1'] / merged_df['count_df2']

办法2：用map映射，更高效 #

不需要全量合并，直接把df2转成Id到count的字典，再映射到df1上计算：

# 生成Id对应count的字典
id_count_map = df2.set_index('Id')['count'].to_dict()
# 映射后计算权重
df1['weights'] = df1['count'] / df1['Id'].map(id_count_map)

这两种方法都能正确算出weights列，不会再报KeyError。

内容的提问来源于stack exchange，提问作者A.N.