Python基于pandas三列数据groupby处理生成对称距离矩阵

完整可运行代码 #

import pandas as pd
import numpy as np

# 原始数据
data = pd.DataFrame({'o':[1,1,2,1,2],
                     'd':[1,2,1,3,3],
                     'distance':[3,5,10,6,8]})

# 1. 计算双向OD对的平均距离
# 生成排序后的临时分组键，保证正向、反向的同组OD对会被分到同一组
data['min_od'] = data[['o', 'd']].min(axis=1)
data['max_od'] = data[['o', 'd']].max(axis=1)
# 分组计算平均值后回填到所有匹配行
data['distance'] = data.groupby(['min_od', 'max_od'])['distance'].transform('mean')
# 删除临时生成的分组列
data = data.drop(columns=['min_od', 'max_od'])

# 2. 生成对称矩阵
dist_pivot = data.pivot_table(index='o', columns='d', values='distance', fill_value=0)
# 对角线统一设置为0
np.fill_diagonal(dist_pivot.values, 0)
# 转numpy数组得到最终结果
res = dist_pivot.values
print(res)

运行输出即为目标结果：

[[0.  7.5 6. ]
 [7.5 0.  8. ]
 [6.  8.  0. ]]

实现逻辑说明 #

• 处理双向OD对平均距离：通过对每行o、d取值做大小排序生成统一分组键，避免(o=1,d=2)和(o=2,d=1)被识别为不同分组，分组计算平均值后回填即可保证双向OD对的距离值一致
• 生成对称矩阵：通过透视表将OD对转成行为o、列为d的表格结构，再用numpy内置方法统一将对角线值设置为0，最终取值即可得到对称距离矩阵

内容的提问来源于stack exchange，提问作者xc-2021