CODEWITHSUNDEEP
pythonpython-3.xnumpyoptimizationsum
Michal
Michal

Reputation: 2073

Map numpy array and sum values on positions in another array

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?

Upvotes: 3

Views: 790

Answers (1)

Divakar
Divakar

Reputation: 221614

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

Upvotes: 4

Related Questions