python - 在另一个数组中的位置映射 numpy 数组和求和值

Question

我有 2 个 numpy 数组:

distances = np.array([
    [0, 6, 7, 4, 2],
    [6, 0, 6, 3, 8],
    [7, 6, 0, 9, 3],
    [4, 3, 9, 0, 4],
    [2, 8, 3, 4, 0]
])

path = np.array([3, 1, 2, 0])

在path , 有一些我想去的地方。每个位置都是 distances 的索引.所以在上面的例子中，我访问了位置 3 , 然后 1 , 然后 2 , 然后 0然后回到3 .为此，我想计算距离总和 3, 1 , 1, 2 , 2, 0和 0, 3 (路径是封闭的圆圈)。

你可以得到3之间的距离和 1来自 distances[3, 1] (= 3)，介于 1 之间和 2来自 distances[1, 2] (= 6) 等。上例中的距离应为 20 .

我创建了这个函数:

import time
import numpy as np

distances = np.array([
    [0, 6, 7, 4, 2],
    [6, 0, 6, 3, 8],
    [7, 6, 0, 9, 3],
    [4, 3, 9, 0, 4],
    [2, 8, 3, 4, 0]
])

path = np.array([3, 1, 2, 0])

def get_distance(path):
    distance = 0

    for i in range(len(path) - 1):
        distance += distances[path[i], path[i + 1]]

    distance += distances[path[-1], path[0]]

    return distance

start = time.time()

for i in range(100000):
    get_distance(path)

end = time.time()
print(end - start) # 0.206

它正在运行，但我需要在长路径(例如 50 个位置)下多次运行此函数(例如百万次)，因此我想对其进行优化。有什么办法吗？

Answer 1

我们可以简单地将索引数组切片并索引到 distances 数组中以获得那些特定的距离，然后将它们相加以获得最终输出。与原始代码一样，我们将单独对最后一对求和。因此，这将是 -

distances[path[:-1], path[1:]].sum() + distances[path[-1], path[0]]