为含图节点数组的Pandas DataFrame添加路径距离计算列

图定义（本问题所用图固定） #

节点连接关系及边距离：
N1 <---10---> N2 <---30---> N3 <---20---> N4

• 节点到自身的距离始终为0
• 任意两节点间的距离为路径上所有边的距离之和

原始Pandas DataFrame #

需求 #

为上述DataFrame添加名为distance的新列，计算每条graph路径的总距离，最终结果如下：

计算示例（以code=06行为例） #

总距离 = (N1 -> N3) + (N3 -> N2) + (N2 -> N4) + (N4 -> N4)
其中：

• N1 -> N3 = (N1 -> N2) + (N2 -> N3) = 10 + 30 = 40
• N3 -> N2 = 30
• N2 -> N4 = (N2 -> N3) + (N3 -> N4) = 30 + 20 = 50
• N4 -> N4 = 0
  总距离 = 40 + 30 + 50 + 0 = 120

实现思路与代码 #

1. 构建节点距离映射表 #

直接用字典存储所有节点对的距离，查询效率更高：

# 节点间距离映射，键为(node1, node2)元组，值为对应距离
node_distances = {
    ("N1", "N1"): 0, ("N2", "N2"): 0, ("N3", "N3"): 0, ("N4", "N4"): 0,
    ("N1", "N2"): 10, ("N2", "N1"): 10,
    ("N1", "N3"): 40, ("N3", "N1"): 40,
    ("N1", "N4"): 60, ("N4", "N1"): 60,
    ("N2", "N3"): 30, ("N3", "N2"): 30,
    ("N2", "N4"): 50, ("N4", "N2"): 50,
    ("N3", "N4"): 20, ("N4", "N3"): 20
}

如果习惯用DataFrame存储邻接矩阵，格式如下：

2. 定义路径总距离计算函数 #

遍历路径中的连续节点对，累加每对节点的距离：

def calculate_total_distance(path):
    total = 0
    # 遍历路径里的每一组连续节点
    for i in range(len(path) - 1):
        node_a = path[i]
        node_b = path[i+1]
        total += node_distances[(node_a, node_b)]
    return total

3. 为DataFrame添加新列 #

调用apply方法处理graph列，生成distance列：

import pandas as pd

# 构造原始DataFrame
data = {
    "code": ["01", "02", "03", "04", "05", "06"],
    "other": ["blue", "red", "green", "white", "blue", "white"],
    "graph": [["N4", "N2", "N2"], ["N1", "N2"], ["N1", "N1"], ["N1", "N3", "N4"], ["N3", "N4", "N1"], ["N1", "N3", "N2", "N4", "N4"]]
}
df = pd.DataFrame(data)

# 添加distance列
df["distance"] = df["graph"].apply(calculate_total_distance)

# 输出结果
print(df)

注意事项 #

• 原始思路中df['graph'].apply(lambda x: get_distance(x['graph']), axis=1)存在错误：对Series调用apply时，每个参数x就是graph列的单个路径列表，无需再取x['graph']；若对DataFrame行调用apply(axis=1)，才需要通过row['graph']提取路径。
• 字典查询比DataFrame查询效率更高，适合数据量较大的场景。

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