Map numpy array and sum values on positions in another array

Question

I have 2 numpy arrays:

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])

In path, there are locations, which I want to visit. Each location is index to distances. So in example above, I visit location 3, then 1, then 2, then 0 and then back to 3. For this I would like to calculate sum of distances 3, 1, 1, 2, 2, 0 and 0, 3 (path is closed circle).

You can get distance between 3 and 1 from distances[3, 1] (= 3), between 1 and 2 from distances[1, 2] (= 6), etc. Distance in example above should be 20.

I created this function:

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

It's working, but I need run this function many times (e. g. milion times) with long path (e. g. 50 locations), so I would like to optimize it. Is there any way?

@iliar Seemed pretty simple, but not a dupe it seems. Posted answer. Thanks. — Divakar, Nov 23 '19 at 09:34

score 4 · Accepted Answer · answered Nov 23 '19 at 09:33

We can simply slice the indexing array and index into the distances array to get those specific distances, then sum those for the final output. As with the original code, we will sum the last pair separately. Hence, it would be -

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

Map numpy array and sum values on positions in another array

1 Answers1

Linked