Home > Back-end >  Fast search for index of element in array
Fast search for index of element in array

Time:08-30

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)

CodePudding user response:

For your prepared example, using argsort will help:

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

print(aa   bb   cc)
  • Related