CODEWITHSUNDEEP
pythonnumpyperformance
Gabriel
Gabriel

Reputation: 42329

Fast search for index of element in array

This is a rather simple problem but I have not found a way to solve it. I need to find the indexes for all the elements in an array, that are stored in different (shuffled) arrays.

I can do that with the code below using the .index() method but this is very slow even for not so large arrays.

Is there a way to improve the performance of this task?

import numpy as np

N = 10000

# Initial array of unique integers
arr = np.arange(0, N)

# Arrays of shuffled elements in arr
np.random.shuffle(arr)
arr1 = np.array(arr)
np.random.shuffle(arr)
arr2 = np.array(arr)
np.random.shuffle(arr)
arr3 = np.array(arr)

# Convert to lists to be able to use the .index() method
arr1 = list(arr1)
arr2 = list(arr2)
arr3 = list(arr3)

# Obtain the sum of the positions in the 3 arrays, for each element in arr
idx_sum = []
for i in range(N):
    i0, i1, i2 = arr1.index(i), arr2.index(i), arr3.index(i)
    idx_sum.append(i0 + i1 + i2)
idx_sum = np.array(idx_sum)

Upvotes: 0

Views: 53

Answers (1)

Ali_Sh
Ali_Sh

Reputation: 2816

For your prepared example, using argsort will help:

aa = np.argsort(arr1)
bb = np.argsort(arr2)
cc = np.argsort(arr3)

print(aa + bb + cc)

Upvotes: 1

Related Questions