时间戳

首页 > 解决方案 > 以矢量化方式查询距离矩阵

python - 以矢量化方式查询距离矩阵

问题描述

一般问题是在给定距离矩阵的情况下，找到列表中每个城市之间的总距离，然后返回第一个城市。

距离矩阵的一个例子。

以下是该问题的有效解决方案，但我不确定如何将其转换为矢量化解决方案。“矩阵”是包含距离矩阵的熊猫数据框，城市是城市列表。

import pandas as pd
def total_distance(matrix, cities):
    dist = matrix[cities[0]][cities[-1]]
    i = 0
    while i < len(cities)-1:
        dist += matrix[cities[i]][cities[i+1]]
        i+=1
    return dist

标签： pythonpandasvectorization

解决方案

如果df是您的距离矩阵，则先将其堆叠。只做一次，因为这是一项昂贵的操作。

dfstacked = df.stack()

接下来，以矢量化方式计算距离：

dfstacked.loc[list(zip(cities, np.roll(cities, -1)))].sum()

MWE：

df
   A  B  C  D
A  0  5  5  3
B  5  0  5  2
C  5  5  0  4
D  3  2  4  0

cities = ['A', 'D', 'B']
dfstacked = df.stack()


dfstacked.loc[list(zip(cities, np.roll(cities, -1)))].sum()
10  # 3 + 2 + 5

推荐阅读