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

Question

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

An example距离矩阵。

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

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

Answer 1

如果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