Searching an array for a value, then printing corresponding values from other arrays/ndarrays of the same length

user_14468532 0
Viewed 58

I have an inefficient algorithm that I need to improve.

Essentially I have arrays:

ids = [134543, ..., 234]
a = [123, ..., 3546]
b = [[435, ..., 549][245, ..., 4986]]

all of the same length (i.e. ids.shape = (600000,) a.shape = (600000,) b.shape = (600000, 2))

and a smaller array of ids of note (length approx 100)

ids_important = [345, ..., 549]

I want to find the indices of the important ids in my ids array, and then output the corresponding elements in a and b.

My current algorithm to do this is:

for i in range(ids_important.shape[0]):
    for j in range(ids.shape[0]):
        if(ids[j] == ids_important[i]):
           print(a[j])
           print(b[j,0])

this algorithm is incredibly slow given the size of the array. I've been told I can improve it using masked arrays but haven't been able to work out how to implement that. Help would be very much appreciated.

numpy-ndarray python-3.x numpy algorithm python
2 Answers

This might not be the most space efficient way of doing it (allocates a len(important_ids) x len(ids) array) but it should be much faster than the other answer and your original approach because it takes advantage of numpy vectorization (and avoids slow loops):

import numpy as np
ids = np.random.randint(0, 15, size=20)
important_ids = np.random.randint(5, 9, size=10)
a = np.random.randint(10, 30, size=50)
b = np.random.randint(10, 30, size=(50, 2))

equals = np.argwhere(ids[None,:] == important_ids[:, None])
inds = equals[:, 1]
print(a[inds])
print(b[tuple(inds), 0])
user_152702870

If you sort the important ids and the arrays, you won't have to iterate over all of the ids every time you search for one.

    ids_important.sort()
    # Sorting an array based on ids of another
    # X = ["a", "b", "c", "d", "e", "f", "g", "h", "i"]
    # Y = [ 0,   1,   1,    0,   1,   2,   2,   0,   1]

    # Z = [x for _,x in sorted(zip(Y,X))]
    # print(Z)  # ["a", "d", "h", "b", "c", "e", "i", "f", "g"]
    a_sorted = [elmt for _,elmt in sorted(zip(ids, a))]
    b_sorted = [[elmt for _,elmt in sorted(zip(ids, b[0]))], [elmt for _,elmt in sorted(zip(ids,b[1]))]]

Then, you can simply iterate over the ids and the important_ones, which will be linear runtime.

i, j = 0, 0
l = len(important_ids)
while i < l and j < 600000:
    if ids[j] < important_ids[i]:
        j += 1
        continue
    if ids[j] == important_ids[i]:
        print(a[j])
        print(b[j, 0])
    i += 1
user_140552590
Related